From c65aa179db7bdd61e2c2821eac87f208a105c053 Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 14 Aug 2026 20:55:37 +0000 Subject: [PATCH 01/18] Open-source X Recommendation Algorithm --- README.md | 17 +- .../src/fixtures/mock_rules_user.yaml | 2 +- bdsm/README.md | 14 +- bdsm/proto/abuse_inference.proto | 3 + bdsm/runtime/score_results_sink_focal.py | 55 +- bdsm/runtime/sink_policy.yaml | 6 + bdsm/tests/test_sink_policy.py | 41 + grox/flows/reply_spam/task_filter.py | 4 +- .../phoenix_candidate_pipeline.rs | 11 + .../filters/brazil_2026_election_filter.rs | 1573 +++++++++++++++++ home-mixer/filters/mod.rs | 1 + home-mixer/params/param.rs | 62 + home-mixer/scorers/author_cold_start.rs | 274 ++- home-mixer/scorers/ranking_scorer.rs | 29 + phoenix-rankall/src/config/mod.rs | 3 +- .../common/xai-recsys/src/model_config.rs | 4 + phoenix/crates/common/xai-recsys/src/util.rs | 196 +- .../serving/xai-recsys-engine/src/python.rs | 3 + .../xai-recsys-engine/xai_recsys_engine.pyi | 2 + .../serving/xai-recsys-proto/src/lib.rs | 2 + phoenix/xrex/models/transformer.py | 1 - 21 files changed, 2271 insertions(+), 32 deletions(-) create mode 100644 home-mixer/filters/brazil_2026_election_filter.rs diff --git a/README.md b/README.md index ca947bf9..fb83f9ca 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) +- [Latest Updates](#latest-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 +## Latest 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. 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/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/bdsm/README.md b/bdsm/README.md index ec45c545..8f072412 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 @@ -110,5 +110,13 @@ The scorer publishes an 8-wide row in `heads.HEAD_NAMES` order. 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. +- 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** (MIN_ACTIONS, 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..818d6103 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): + """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 30 actions. 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 < 30: + 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): @@ -1116,7 +1165,7 @@ 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), "labels": user_labels, "model_version": args.model_version, "pipeline_version": "gpu_scorer_kafka_v3", diff --git a/bdsm/runtime/sink_policy.yaml b/bdsm/runtime/sink_policy.yaml index eda36890..c87e9be7 100644 --- a/bdsm/runtime/sink_policy.yaml +++ b/bdsm/runtime/sink_policy.yaml @@ -2,6 +2,12 @@ version: "2026-08-13" # Public-release note: the per-head operating points below (the 2-D +# [tau, lambda] pairs, cusp_delta, and reply_spam_hard_suspend_tau) are +# REDACTED in the open-source export — they are replaced with an +# out-of-range 9.99 sentinel (these are probabilities in [0, 1], so a +# 9.99 threshold never fires and is plainly a stub, not a real value). +# 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 diff --git a/bdsm/tests/test_sink_policy.py b/bdsm/tests/test_sink_policy.py index 89db1c5c..475346d9 100644 --- a/bdsm/tests/test_sink_policy.py +++ b/bdsm/tests/test_sink_policy.py @@ -43,3 +43,44 @@ 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() + 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) is None + assert m.build_enforcement_note(heads, None) is None + assert m.build_enforcement_note(heads, []) is None + assert m.build_enforcement_note(_heads(FollowBot=0.4, LegitimateUser=0.8), long) is None + + note = m.build_enforcement_note(heads, long) + 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/flows/reply_spam/task_filter.py b/grox/flows/reply_spam/task_filter.py index 7b39d7d2..3bf626a8 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 = 30000 @override @classmethod @@ -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 = 30000 @override @classmethod diff --git a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index dd38792c..3595411e 100644 --- a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs +++ b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs @@ -37,6 +37,7 @@ 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; @@ -355,6 +356,16 @@ impl PhoenixCandidatePipeline { Box::new(PreviouslyServedPostsFilter), Box::new(MutedKeywordFilter::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), 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..2f832ca7 --- /dev/null +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -0,0 +1,1573 @@ +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 + +// 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. + +// OmarAzizSenador deleted his account at the time this code was written. + +/// User ids reported to the Electoral Court for the Brazil 2026 election. +static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(|| { + FxHashSet::from_iter([ + // @renildo + 14160928, + // @renatoroseno + 14492205, + // @pedro_lupion + 15022409, + // @Sen_Cristovam + 20242549, + // @marcelvanhattem + 21069302, + // @nilvanferreira + 21571098, + // @RafaellMilas + 21857669, + // @chagasvieira + 22035789, + // @GabrielSouza_RS + 22480147, + // @Pimenta13Br + 22864100, + // @rigotto + 23443097, + // @euserafimcorrea + 24761475, + // @ManuelaDavila + 25858078, + // @CintyaMuniz + 26284058, + // @depguibismarck + 29026134, + // @cirogomes + 33374761, + // @aldenorlima + 33973761, + // @jooliveirapb + 34909888, + // @caduxavier + 35470350, + // @NetoAM + 36145775, + // @marciokieller + 36403221, + // @AlicePortugal + 36971658, + // @eniobritodesa + 37654300, + // @dep_geraldo + 37711911, + // @inacioarruda + 38659819, + // @murilogaldinopb + 40040727, + // @FlavioBolsonaro + 40053694, + // @BetoAlbuquerque + 40463380, + // @alexandrekalil + 40721032, + // @jessicamichels + 41355867, + // @lcbusato + 42300711, + // @Alice_Portugal + 42304398, + // @DepArthurMaia + 42454330, + // @kruke1 + 42487937, + // @marcelorangel1 + 44107208, + // @perpetua_acre + 44693900, + // @PedroTaquesMT + 45448098, + // @betinhogomes + 46958570, + // @DanielVilela15 + 47371488, + // @zeca_dirceu + 47461491, + // @eduardopaes + 48298703, + // @RicardoBarrosPP + 50101324, + // @danielxdonizet + 50360881, + // @anyortiz + 50430144, + // @priscilakrause + 51066167, + // @advcorrea + 51169114, + // @afranioboppre50 + 51735736, + // @DeputadoBacelar + 52364013, + // @JuniorMochi + 52483984, + // @GuilhermePaz + 52750366, + // @rosanedopv + 52842057, + // @RodrigoGuedesam + 52954632, + // @RollembergPSB + 53050115, + // @jorginhomello + 53073647, + // @nelsinhotrad + 54412355, + // @RubensOtoni + 54539654, + // @DeputadoWelter + 54545619, + // @marcosdoval + 54600557, + // @mineiroptrn + 54897084, + // @dep_acoutinho + 56480030, + // @venezianovital + 56734024, + // @DelegadoMeneses + 57108807, + // @RobinsonFaria + 57151676, + // @HarifeViegas + 57163107, + // @AlineMariano_pe + 58314645, + // @fredlinhares + 58328715, + // @JutayMeneses + 59243227, + // @fernandofilhope + 59323812, + // @franzepiaui + 61190865, + // @ale_campelo + 61325857, + // @carlaayres + 62286707, + // @HenriqueFontana + 62804559, + // @BohnGass + 63127680, + // @JoseAirtonPT + 63868587, + // @Bobadra + 64302060, + // @Marco_Brasil + 64434503, + // @livioluciano + 65059970, + // @FatimaCleidePT + 65491379, + // @fabio_novo + 66413485, + // @miguelcoelhope + 66525428, + // @VerGuilherme + 66704302, + // @cantojocelito + 66719730, + // @betodoisaum + 66753261, + // @mateus_simoesmg + 66789220, + // @ciro_nogueira + 67098726, + // @PattyParente + 67234932, + // @acirgurgacz + 67601773, + // @eduardoreiner + 67947379, + // @andrekamai + 68404402, + // @costa_rui + 68466700, + // @MarcosSantanaSC + 68554516, + // @CarlosBolsonaro + 68712576, + // @ProfIsrael + 68719944, + // @juliophilbert + 69073488, + // @Gyselle_Soares + 70131422, + // @ernanipolo + 70284501, + // @deputadoismael + 70453020, + // @FaissalCalil + 71602909, + // @SamuelMalafaia + 73442547, + // @pauderney + 73803827, + // @geraldoalckmin + 74215006, + // @eduardobismarck + 74361905, + // @luizcoutopt + 74738674, + // @agenorsantospa + 74867163, + // @santinroveda + 75849275, + // @fabriciopref + 76039110, + // @leitaothales + 76224437, + // @DepMajorAraujo + 76383384, + // @DepAfonsoHamm + 76741399, + // @MarcelAlexandre + 77266759, + // @raquelferreirar + 78707944, + // @joslene65 + 78714361, + // @profdorinha + 79174387, + // @capitaotadeu + 80214491, + // @carlosmatosce + 80597262, + // @depdarcidematos + 80712669, + // @Adrianageronim + 80733107, + // @HermanoMorais + 80995814, + // @fabiofelixdf + 81384580, + // @liviaduartepsol + 81742151, + // @carlosjordy + 82271629, + // @Manato_es + 82433790, + // @darypagung + 82936046, + // @joaoromaneto + 84141432, + // @andrerochamg + 84196661, + // @wevertonrocha + 85859074, + // @realrcoutinho + 86065271, + // @vozdaenf + 86348991, + // @romeropelapb + 88303403, + // @arthurgurgeladv + 90630534, + // @ANDREAESPANHA + 90898890, + // @wilsonlimaAM + 91109801, + // @f_trad + 92509126, + // @marcelodeputado + 93116173, + // @tomazteixeira + 93958167, + // @coroneldavidms + 94378206, + // @ruycarneiropb + 94428241, + // @samanthacavalca + 94806323, + // @requiaooficial + 95253000, + // @cabogilberto + 95526088, + // @anapaulalimapt + 95939603, + // @PROFTULIO + 96570084, + // @maxmacieldf + 98786988, + // @michelyfarina + 99617723, + // @julioarcoverde + 101016613, + // @PabloValenteDF + 102257530, + // @NicolasTrancho + 106110720, + // @DivaneidePT + 107976868, + // @mayradiasam + 108527180, + // @HugoMottaPB + 108988113, + // @deplucasdelima + 109318072, + // @adjutoafonso + 109332428, + // @joaopaulodopt + 109657263, + // @HendersonPinto + 111717686, + // @alielmachado + 115519533, + // @luizaogoulart + 116541810, + // @adelmosoaress + 116751115, + // @acrisiosena + 117425043, + // @vereadorjulio + 117594353, + // @julio_cesar_pi + 119115954, + // @adrianogaldino + 119586707, + // @sibellebarros + 120535860, + // @aryvanazzi + 121422504, + // @walteralvesrn + 121425571, + // @netoevangelista + 121594926, + // @felipecarreras + 122184686, + // @delegadowaldir + 123689660, + // @MarceloCastroPI + 125822264, + // @jardelinacioam + 126323536, + // @deputadopezenti + 127708617, + // @ricardozguidi + 127975238, + // @Mersinho_Lucena + 129681031, + // @alexandrebaldy + 130620293, + // @depdelmasso + 132190299, + // @PedrodoOvo + 132220469, + // @Larissafgaspar + 132480664, + // @glaustindafokus + 133692726, + // @sidneyleite_ + 135349877, + // @LucianoDucci + 136004062, + // @tadeuveneri + 137919701, + // @MerlongSolano + 140413929, + // @lucianogenesio + 142068227, + // @silvioantonioma + 142501634, + // @Carloshbfavaro + 143529694, + // @vicentinhojr + 149013361, + // @pedrogomatos + 152900822, + // @dacassia1 + 155768056, + // @profcanguru + 156326262, + // @AzambujaReinald + 157184730, + // @MarcosRogerio + 160895960, + // @MarcoVamosaLuta + 161318399, + // @marinorpsol + 161404367, + // @Vanderlan_VC + 162807255, + // @TJMFernandes + 163251815, + // @EduardoBraga_AM + 164439493, + // @crisrbritto + 165329128, + // @Jeronimoba13 + 165754961, + // @DeputadoRoberio + 165914831, + // @andrepdt12 + 165922194, + // @geovaniadesa + 165936805, + // @caiopontess + 165949065, + // @NelterQueiroz + 166164275, + // @deputadotoninho + 166417133, + // @marcocastilhoss + 168408841, + // @gustavohaguera + 168653875, + // @TiberioLimeira + 168657354, + // @walterlfcaval + 170638771, + // @JanineLucena + 172579844, + // @tayannystefany + 172581180, + // @toinharocha2 + 173823248, + // @raquellyra + 174370381, + // @CinthiaCRibeiro + 175438359, + // @queirozmfilho + 175901537, + // @alxlindenmeyer + 179749078, + // @VictorCoelhoES + 182050588, + // @rodfvale + 186331377, + // @MariaSeffair + 199526953, + // @Anderson_Prego + 201100473, + // @natbonavides + 203332874, + // @AntonioCoelhoPe + 208636149, + // @MeuEuPolitico + 209674215, + // @saullovianna + 210299199, + // @fatimacnagitos + 212950503, + // @celmarcosantos + 219262636, + // @arilsonchiorato + 220353446, + // @boscosaraiva + 221399756, + // @capitaowelton + 222230527, + // @gardelrolim + 226984087, + // @osmarfilhoma + 230495542, + // @luisa_canziani + 230950274, + // @deltapericles + 239036359, + // @silascamara_ + 243018634, + // @UrsulaVidalPA + 243429150, + // @DeAssisDiniz + 247906787, + // @ranypaulino + 252553750, + // @rdnarede + 254172269, + // @deprsantos + 254201392, + // @wellington_luiz + 255637975, + // @macielchris + 261472149, + // @coronelfrota + 263153198, + // @catulejr + 264391266, + // @alinegurgel_ap + 267981392, + // @ayres_jr + 269938072, + // @alexgalvaodf + 272477039, + // @diegogarciapr + 273616279, + // @doutorgutemberg + 274515672, + // @marcelomaranata + 287876093, + // @delegado_waldir + 288735634, + // @coelho_rodrigo + 290204659, + // @maitebrusman + 294065810, + // @depkleberRN + 302056921, + // @WedersonLopes + 313684933, + // @_sergiosouza + 317973567, + // @danielvalencapt + 318019551, + // @depmaracaseiro + 321414691, + // @loureirocris + 322406139, + // @jorgeviana + 325131009, + // @simboramudar22 + 337269106, + // @marisa_lobo + 340331807, + // @MarceloBelinati + 345512946, + // @edusantosdf + 348089005, + // @leilafonsecapb + 349053302, + // @paulinhoramosap + 349163261, + // @andrefernm + 360231929, + // @milkleileite + 365210575, + // @andresalineiro + 367519089, + // @Ana_claudiapb + 370942708, + // @JorgeFrederico2 + 398245257, + // @brena_dianna + 412508127, + // @faf_freitas + 412669574, + // @zeliomota + 420614002, + // @vitordeangelo + 424455116, + // @instrutormarcio + 437397340, + // @AllysonBezerra_ + 444925483, + // @CarlosMoises + 456795842, + // @JulioCesarRib + 459681629, + // @chicorafa_s + 487108193, + // @pedroluislongo + 487622592, + // @MoisesSantosAc + 538269241, + // @JulianaBRedivo + 545696318, + // @alcyvania + 556852639, + // @OthelinoNeto + 583377940, + // @AfonsoFlorence + 599427558, + // @elmanooficial + 626602522, + // @fabiogov55 + 630758039, + // @helenaduailibe_ + 632737286, + // @mitchellemeira + 746090918, + // @suelenmarques06 + 796882578, + // @JenirNeves + 813802178, + // @carlaopelobem + 893975196, + // @Isoldadantaspt + 999393290, + // @SHEILAKLENER + 1038802238, + // @rodrigostm10 + 1038849745, + // @xambinhoes + 1068257582, + // @tatianehelena81 + 1075326786, + // @depjanetedesa + 1094959356, + // @prof_juniorgeo + 1226451780, + // @marcoswesleymw + 1308817182, + // @alexceoficial + 1314840228, + // @D_GoretePereira + 1362596354, + // @mickasevalho + 1461892208, + // @deyvidbacelar + 1526183576, + // @NegrahLima + 1571332381, + // @marciopachecopf + 1613063005, + // @GilvanMaximoOfc + 1651852124, + // @Brandaveneno + 1710385291, + // @dr_furlan + 1844887189, + // @valmirdesergipe + 1960681207, + // @helinhocastro + 1977089148, + // @DepAntoniaLucia + 2161527577, + // @LUANARUIZSILVA + 2216639570, + // @DepFederalMoses + 2217650233, + // @BrunoCarianha + 2299610603, + // @DavidAlmeidaAM + 2353403137, + // @capitaoassis10 + 2359974485, + // @Francadf_ + 2445938091, + // @Alfredoficial22 + 2474258532, + // @yuriarrudam + 2495584641, + // @prof_rosaneide + 2523520542, + // @viagensdaiw + 2572998767, + // @anadogasoficial + 2583179096, + // @barbosinhams + 2612421128, + // @EderMauroPA + 2632802395, + // @paulo_litro + 2647646724, + // @LuizianneLinsCE + 2666819714, + // @LulaOficial + 2670726740, + // @marcoaurelioITZ + 2691147288, + // @LucasVergilioGO + 2717062461, + // @netto_expedito + 2717313965, + // @tiaomedeiros + 2732504341, + // @mqueiroga22 + 2823869072, + // @AlbertoMaiaf + 2834745257, + // @paulolemosap + 2881293743, + // @manascimentogo + 2892653500, + // @LizianyM + 2894163148, + // @carlosveraspt + 3010698441, + // @mendanhagustavo + 3018780587, + // @karlosbernardoA + 3029604164, + // @zuleidequeirozf + 3130842411, + // @meire_cruvinel + 3205786257, + // @Marcio_Honaiser + 3294107902, + // @deputadopriante + 3305760711, + // @moisesbrazpt + 3373574517, + // @kleybe_morais + 3512854216, + // @carmelonetobr + 3662771592, + // @paulo_mavignier + 4043244676, + // @zeaugustonalin + 4056653428, + // @_AlessandroSE + 4250596815, + // @CarolDeToni + 4566967516, + // @FofaBorges + 4775264669, + // @AndreaOficial55 + 703604819538407424, + // @MARTACLERIALIMA + 713343399898820608, + // @LiaGomesCE + 724003686863740933, + // @ViniciusFerroC + 726243481669242880, + // @brisabracchi13 + 728292397130592256, + // @cabo_senna + 744609688415789056, + // @jarir_pereira + 746804180430487554, + // @angelaamin11 + 766624608338477056, + // @vanessapstubh + 770298023582765057, + // @Jorgepinheiroof + 807903184576532480, + // @amorimvivian_ + 820867290749161472, + // @ProfessorEuler + 824285052590845955, + // @angelocoronel_ + 829391158208032768, + // @AdelitaMonteiro + 830144764360192002, + // @ranallipf + 832322309436403712, + // @GeneralGirao + 841700087143288832, + // @WillaceSouza + 850775334362505216, + // @vivitobiasms + 852960064159842305, + // @rodrigodiasbsb + 863812402680397826, + // @pedroalmeidace + 882539177757298689, + // @meuamigojoao + 899602419302240257, + // @todandara + 904842960931565568, + // @AlcyPinheiroCE + 909049597091356673, + // @TorminCassiana + 931087564064329728, + // @rigoni_felipe + 939423395376136192, + // @sergiokruke + 940210949943910401, + // @WilliamSiriRJ + 958357145761845248, + // @dep_paulinha + 973982865997385728, + // @Diegoferdfc + 977195483507707904, + // @ZeCocaOficial + 978602905690427392, + // @Renatoafjr + 982217430297554945, + // @fernandomaximoX + 983360917227364353, + // @RenilceOficial + 984221723544444928, + // @eng_angelo44 + 984522534623301632, + // @IndiaraNOVO + 986396254287646721, + // @leaolais_ + 986659343436255232, + // @ptmatiaspedro + 986903953014128641, + // @KerexuOficial + 992801222926196736, + // @kekabagno + 993288307625943040, + // @annakarinapsol + 1009608950176796677, + // @CapitaoContar + 1011667518728089602, + // @Trom_Petista + 1013494732834639874, + // @limmapiaui + 1014568888078651392, + // @acaroldartora + 1016697971843502080, + // @RomeuZema + 1020798087449776128, + // @EduGiraoOficial + 1024403315164160000, + // @Laurez_Moreira + 1025352673292378113, + // @zenaidern + 1025870063461720064, + // @BocaAbertaOf + 1031914738933018624, + // @RafagninLuciana + 1034079788162535424, + // @erikaamorimce + 1034138653147242496, + // @AmauriRibeiroGO + 1036681658760658945, + // @RenanSantosMBL + 1042601099566436352, + // @raphaelbarrabr + 1049870508643233793, + // @pluviapt + 1062505159824150530, + // @capitaocarpe + 1062678892216020992, + // @clarissatercio + 1065379080084865024, + // @capalbertoneto + 1076135802969772032, + // @tarcisiogdf + 1078618844007157761, + // @MarcosA_Sampaio + 1080911369447399430, + // @capitao_alden + 1081245039920144384, + // @Khalill_gui + 1083098735675084800, + // @wilkerbarretoam + 1085537170276958208, + // @PlinioValerio45 + 1091691201844207617, + // @doriel_barros + 1092446595889745921, + // @joaoluizam + 1092549114762539009, + // @betopereirams + 1092816222012534784, + // @biologiagabriel + 1092879030465032192, + // @ZequinhaMarinho + 1093223195258298368, + // @DepGuiLandim + 1095422600854032390, + // @MondardoGiovana + 1095645491419852800, + // @JuniorManoDep + 1095696286974652422, + // @IdilvanAlencar + 1095754303367720960, + // @depsargentolima + 1095990374550695938, + // @DFDanielFreitas + 1097498693719199744, + // @CelAlfredo + 1098328459825287187, + // @benesleocadiorn + 1099003656064643073, + // @depPedroLucasF + 1104120809482866688, + // @PollonMarcos + 1105081277135417345, + // @CoelhoAntonioPE + 1106161642486796288, + // @SargentoBetania + 1107097837194694663, + // @JoaoCampos + 1112804630453501953, + // @SF_Moro + 1113094855281008641, + // @faustojr_am + 1116362076459601920, + // @10ronaldomartin + 1118164353121968129, + // @deputadaniella + 1118481835376431104, + // @flavionpi + 1118489062866849793, + // @henriquecesarr + 1121117485900730368, + // @georgebastos30 + 1141394024860966912, + // @CarboniRogerio + 1143607219302387712, + // @pedrorochafilho + 1151964529896644620, + // @washingtonban + 1154733508088270849, + // @karlasarney_ + 1160505413697253376, + // @jofariasm_ + 1167172974451073024, + // @RobertoCidadeAm + 1168549394830020608, + // @jadearomero + 1171411853282557952, + // @prjuniortercio + 1171959927771975680, + // @FranciscodoPT13 + 1188191474636206082, + // @SandraLimadeVa1 + 1199740816018853890, + // @LeoSuricate + 1208544960032727040, + // @ManuVieiraSC + 1210296676520513536, + // @CruzOrleans + 1211689081861623808, + // @OficialNenemAl + 1216746154626625536, + // @JairSoutoAM + 1218904655591243777, + // @JohnRobertPA + 1219003418347483136, + // @DrVictorAmoras + 1224505061558161408, + // @KlesleyGarcia + 1224615230195609601, + // @profterezinhaPT + 1225158156545904642, + // @brenogaribalde + 1225774842735202306, + // @gerlanebaccarin + 1225861584972664833, + // @deputadoelizeu + 1226865113635926021, + // @dramayraoficial + 1226874382938787841, + // @BabaTupinamba + 1232082766071767045, + // @MassamiMiki + 1234278834519920642, + // @juizcubas + 1234488841479884801, + // @JacqueMoraes_es + 1235371041712726022, + // @julinhodeputado + 1238202050745446401, + // @eusouamom + 1238868377675980803, + // @fabiolopes_38 + 1241821464111853570, + // @elianexunakalo + 1241878280841674752, + // @samaramartinsup + 1244783574676627460, + // @drluizovando + 1247287375803355136, + // @CoronelFernand9 + 1249769581725573121, + // @jaziel_dr + 1251558672821600259, + // @cirineu_costa + 1251676296842805257, + // @fcordeirosc + 1252561439686037506, + // @fabio_schiochet + 1253698500401008641, + // @tyago_hoffmann + 1254019417148661760, + // @peritatirotti + 1254073103661088768, + // @tanielmacedo + 1254084527422689282, + // @beckhausersc + 1257020913893195779, + // @DraSilvana2 + 1257476064219140097, + // @najuliaribeiro + 1259624231765184514, + // @adailfilhoam + 1260696151029940225, + // @camilajarams + 1261423487585005575, + // @mariadoscamelos + 1267127677468753920, + // @manupeloES + 1269632201760690176, + // @juliermesenav + 1271216272148176896, + // @carladicksonrn + 1273322623380983815, + // @MucioBotelho + 1275549634757365760, + // @LucasCaregnato + 1275636866075803648, + // @victorcarvamt + 1280120303818072064, + // @JuniorGeraldo_ + 1280224335307907077, + // @KodamaThiago + 1281980054730407936, + // @eugenialima_pe + 1285282863471001601, + // @daniportelape + 1285295088525082631, + // @ericodonovo + 1287396715331452932, + // @MatheusLaiola + 1292497353166004227, + // @DepCoronel + 1295725787790942208, + // @gustavosefer + 1298369792169127939, + // @dinhodowsley + 1298603065457680385, + // @luladafonte + 1300279136620085249, + // @marcelinhoguima + 1303443223663325185, + // @chaves_hildon + 1304134002547331072, + // @GuguSeba + 1310253107310460933, + // @CoronelRomualdo + 1310744586139119618, + // @EduardoRiedel_ + 1310977778574032900, + // @PabloSilvaLira + 1311620519935045632, + // @cassymonteiro + 1312841615111843848, + // @delegadkatarina + 1314575606487670784, + // @anaportelams + 1319825203313184768, + // @DeoliAnderson + 1322213572953493505, + // @bocalomoficial + 1335685798746845190, + // @rosanadasaude_ + 1351505730323558400, + // @dimasfabianomg + 1355216660068761606, + // @AnaPimentelmg + 1356434859934298115, + // @faustinorn01 + 1356762322489004032, + // @RafaelFonteles_ + 1358739792163389440, + // @FelipeAlecrimPE + 1375498458132537348, + // @celiaxakriaba + 1379025327314329608, + // @marleipr + 1380366998488645636, + // @apropriajulia + 1382550118314999812, + // @PrefeitoCesar + 1383362690584768514, + // @santanna_cs + 1385208512955957250, + // @sauloportolivei + 1409385758838906882, + // @juniorferrari55 + 1409990621536919555, + // @deputadodrhugo + 1413128451154866185, + // @depclaudiac + 1423660842239791115, + // @coronelbonates + 1427816945735319554, + // @robertadahorta + 1434501105241792512, + // @NFgoes + 1435915606541410305, + // @ScalcoDarlan + 1437758990172184577, + // @joelrodriguespi + 1447921024117493761, + // @atilaliraof + 1448414094441291777, + // @LuizEduardo_RN + 1452666475844706306, + // @josecam01577970 + 1459681826700673030, + // @RicardoArrruda + 1462770465068437504, + // @yurydoparedao + 1478380311532736512, + // @Anderson_ma123 + 1478686290270896128, + // @firmo_oficial + 1481351815669112840, + // @depjaqueline + 1488866652037029889, + // @GayerGus + 1489108473027739654, + // @eribertomfilho + 1491176697445683205, + // @RosaRezendeGO + 1491387105913810944, + // @neyamorimac + 1492921263299379207, + // @sgtgoncalves22 + 1504105286876991489, + // @SKaripuna + 1506313161921777676, + // @orleansbrandao_ + 1510790014111784967, + // @DelRodrigoSa + 1512105284558282765, + // @reillerlopes + 1513329845551390720, + // @franze_carneiro + 1513903701483790343, + // @mariaarraespe + 1515390927958908928, + // @SocorroWaquiim + 1516019212145287168, + // @DanielleDVale + 1516121393028644869, + // @luizgastaoce + 1519011181574434816, + // @drviniciuspi + 1519744542173536259, + // @coronelamadeu + 1523726346232414208, + // @SindoleyMorais + 1525834800233336832, + // @JadyelAlencarr + 1531271393219837956, + // @eudonaneuma + 1533265972659990528, + // @melcafariaspb + 1534926600743051266, + // @rejanepstu + 1539038049937547269, + // @AnneMarques_AP + 1539992525980815364, + // @soududaramos + 1542878687120482306, + // @victorlinhalis + 1543943524189720576, + // @LuisaCela87 + 1545437399446077440, + // @MissiasDias + 1553167499901943808, + // @catanhopt + 1555559206635315201, + // @VitormoreiraCG + 1557086339740405761, + // @Clecio_Luis + 1560635958994812928, + // @DeboraMenezes22 + 1564101654496133120, + // @weibetapeba + 1575549757900132379, + // @Arnaldodeputado + 1578127734362046489, + // @antoniodoidoofc + 1588344483766321154, + // @gianninogueira2 + 1589336373101723649, + // @Marisaloboreal + 1589768586096148483, + // @izaarrudape + 1593200340777803776, + // @davivalencambl + 1594709250835709957, + // @WickRyanAM + 1612178912947183620, + // @deputadairacema + 1612229099551858688, + // @aveiltonsouza + 1612484459034550273, + // @ToGomes5 + 1613522130368356352, + // @tadeudesouzaam + 1618970134005121025, + // @gracinhamaosant + 1624490961236631553, + // @EdsonKambeba + 1633100825403826177, + // @AdrianaLeal2507 + 1637963109573828608, + // @adrianalmeidapt + 1638181881127612416, + // @wistongomess + 1648046380106104833, + // @drbrunoresende_ + 1661373663256408066, + // @DouglasRuas_RJ + 1663563614954086401, + // @bemoreiradf + 1666457401816391682, + // @babatupinamba_ + 1691097420166254592, + // @PauloAssun68233 + 1702440269138866176, + // @hebertcsgyn + 1713936872370532352, + // @moreiramissao + 1721592893511483392, + // @eduardowilliamm + 1728271638104358912, + // @bellaccarmelo + 1743284668764463105, + // @jotinhapiaui + 1744704300616417280, + // @andrabianchessi + 1752837451175907328, + // @Lenesilllva + 1771346241449865216, + // @jorgemaciel05 + 1775273264774045697, + // @CamillaGonda + 1781485235047174144, + // @aucilene_a75929 + 1784209048671322112, + // @eliabecamposs + 1784275708606320641, + // @oescobarpro + 1786921594390016001, + // @coronellrosses + 1793428861348167680, + // @gabi_bvnt + 1809421300433317892, + // @passinhoisa + 1812967972358569984, + // @SandroOmar48237 + 1816185635578978304, + // @007Douggomes + 1821931832008392704, + // @cironogueirapi + 1844442355421614089, + // @DragUrbana + 1847431729415397376, + // @abreu_de63729 + 1855074530613403648, + // @sandsonmenezes + 1856816774517268480, + // @depcabomacielam + 1859357204278542336, + // @Lukaovereador + 1874897810753208320, + // @oiurecastro + 1878758512639033344, + // @_marquinhostrad + 1883938805306052608, + // @bolsonaro__jr + 1884423273490108416, + // @ameliocayresdep + 1887508211445702656, + // @FelipeVasquesce + 1889410539639517184, + // @Rpachecopinho + 1898684990428151808, + // @viscontioficial + 1898913625697292288, + // @GersonClaroMS + 1899113980544897024, + // @glenioseixas + 1909726014202171392, + // @prof_elson_sc + 1937980838299504644, + // @jotabrandaoam + 1949664156606459905, + // @NoronhaPro7153 + 1958622002719182848, + // @RubensAngiolett + 1959013340661035008, + // @Drfilipecm + 1972344363654037504, + // @DrManuelMarcos + 1975317845962858501, + // @CarlosCostaPE10 + 1975873882654441472, + // @marciaabrahaodf + 1985760866109964288, + // @DaClaudio33805 + 1988208205575712772, + // @vandawitoto + 1988323381792681984, + // @SamaraMadureira + 1990473255539621889, + // @profraydf + 2010751427623497728, + // @brenomacedopi + 2013660202873032707, + // @robson_cacau + 2015964533689323520, + // @sicchar1389 + 2018445464908042240, + // @GreguyLoooban + 2027415714953396224, + // @AraceliLemosOF + 2028870149546373120, + // @missaogabi + 2036817223558234112, + // @owilsonmartins + 2043714770591735809, + // @RitaDamore11 + 2047488478372368384, + // @jaimeverruckms + 2052755712644730887, + // @ThefiAmancio21 + 2054040940885532672, + // @onenencoelho + 2054729517864771584, + // @edmartresoitao + 2056559848750276608, + // @obrasilcomlula + 2058595300873289728, + // @vivianeluizams + 2059726216119107587, + // @ProfWiterNaves + 2059790203091341316, + // @DrCrisVeloso + 2061889166779019264, + // @DaversonMatos + 2062692988606717952, + // @HVilelaporGoias + 2069828777517973504, + // @angelaabukce + 2071949683736354816, + // @drthalescoelho + 2075213887213797376, + // @Nayladasilva0 + 2078543961333850112, + // @pituxosergipe + 2080630684989730816, + // @Irmamabelmelo + 2080632627887734784, + // @thiagocampeloof + 2081154349267345408, + // @Profsamuelsiebr + 2084048343609540608, + // @AuriJuniorr13 + 2084644771927076864, + // @CarmemOliverofc + 2084756136004083712, + // @nandovianapsol + 2086122224126136320, + ]) +}); + +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(), 665); + } + + #[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/mod.rs b/home-mixer/filters/mod.rs index b9f09b9a..9cb53f72 100644 --- a/home-mixer/filters/mod.rs +++ b/home-mixer/filters/mod.rs @@ -2,6 +2,7 @@ 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; diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index dc3a12e4..8b0054e4 100644 --- a/home-mixer/params/param.rs +++ b/home-mixer/params/param.rs @@ -279,6 +279,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!( @@ -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", + 5 +); +param!( + ColdStartImpressionScale, + f64, + "rust_home_mixer_cold_start_impression_scale", + 1.0 +); +param!( + ColdStartTrackedIds, + String, + "rust_home_mixer_cold_start_tracked_ids", + "" +); param!( EnableCachedPosts, diff --git a/home-mixer/scorers/author_cold_start.rs b/home-mixer/scorers/author_cold_start.rs index 246d26ce..e7652f85 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; @@ -79,6 +83,54 @@ 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); + if tracked.contains(&c.tweet_id) { + *counts.entry((c.tweet_id, source)).or_insert(0) += 1; + } + if tracked.contains(&c.author_id) { + *counts.entry((c.author_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) } @@ -153,6 +205,55 @@ 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.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, candidates: &[PostCandidate], @@ -164,11 +265,12 @@ fn apply_cold_start( 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 use_ts = query.params.get(EnableColdStartThompsonSampling); let (positions, nonzero) = positions_among_nonzero(scores); let max_cold_start_slot = (query.params.get(LowImpressionsMaxPositionRatio) * nonzero as f64) as usize; - let best = candidates + let eligible: Vec = candidates .iter() .enumerate() .filter(|(i, c)| { @@ -179,9 +281,24 @@ fn apply_cold_start( && c.view_count.is_some_and(|imp| imp < threshold) }) .map(|(i, _)| i) - .max_by(|&i, &j| scores[i].total_cmp(&scores[j])); + .collect(); + + let best_idx = if use_ts { + pick_thompson( + &eligible, + candidates, + scores, + query.params.get(ColdStartBetaAlpha0), + query.params.get(ColdStartBetaBeta0), + query.params.get(ColdStartImpressionScale), + query.params.get(ColdStartTsTopK) as usize, + &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,6 +320,9 @@ impl AuthorColdStart { candidates: &[PostCandidate], scores: &[f64], ) -> Vec { + let tracked_ids: String = query.params.get(ColdStartTrackedIds); + record_tracked_ids(candidates, &tracked_ids); + if !query.params.get(EnableViewerColdStart) { return scores.to_vec(); } @@ -226,6 +346,7 @@ 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, @@ -313,11 +434,21 @@ rust_home_mixer: } fn cold_start_candidate(author_id: u64, age: Duration, view_count: u64) -> PostCandidate { + cold_start_candidate_with_favs(author_id, age, view_count, 0) + } + + fn cold_start_candidate_with_favs( + author_id: u64, + age: Duration, + view_count: 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), + fav_count: Some(fav_count), ..Default::default() } } @@ -360,6 +491,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 } @@ -491,4 +648,111 @@ 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![ + cold_start_candidate_with_favs(1, minutes(10), 0, 0), + cold_start_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![ + 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(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]); + } + + #[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(&(10, pb::ServedType::ForYouPhoenixRetrieval as i32)), + Some(&1) + ); + assert_eq!( + counts.get(&(10, pb::ServedType::ForYouPhoenixRetrievalMoe as i32)), + Some(&1) + ); + assert_eq!( + counts.get(&(20, pb::ServedType::ForYouInNetwork as i32)), + Some(&1) + ); + assert_eq!(counts.get(&(30, 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/ranking_scorer.rs b/home-mixer/scorers/ranking_scorer.rs index 6dfcb7aa..6dc258ac 100644 --- a/home-mixer/scorers/ranking_scorer.rs +++ b/home-mixer/scorers/ranking_scorer.rs @@ -415,6 +415,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 } diff --git a/phoenix-rankall/src/config/mod.rs b/phoenix-rankall/src/config/mod.rs index bbf582d6..d1f8db77 100644 --- a/phoenix-rankall/src/config/mod.rs +++ b/phoenix-rankall/src/config/mod.rs @@ -177,6 +177,7 @@ impl PipelineKind { WindowConfig::new("1fav", 24), 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), @@ -352,7 +353,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/crates/common/xai-recsys/src/model_config.rs b/phoenix/crates/common/xai-recsys/src/model_config.rs index 853aa05d..262a700f 100644 --- a/phoenix/crates/common/xai-recsys/src/model_config.rs +++ b/phoenix/crates/common/xai-recsys/src/model_config.rs @@ -37,6 +37,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 +251,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 { @@ -281,6 +284,7 @@ impl ModelConfig { num_post_bool_features, num_post_float_features, num_post_int64_features, + enable_stale_post, }, history_seq_len, candidate_seq_len, diff --git a/phoenix/crates/common/xai-recsys/src/util.rs b/phoenix/crates/common/xai-recsys/src/util.rs index 2aeeed6f..4ace24eb 100644 --- a/phoenix/crates/common/xai-recsys/src/util.rs +++ b/phoenix/crates/common/xai-recsys/src/util.rs @@ -35,6 +35,10 @@ fn record_sid_coverage(sequence: &str, present: u64, count: u64) { .inc_by(count.saturating_sub(present)); } +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 +47,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 +138,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], @@ -443,22 +456,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 +513,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 +534,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 +560,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,6 +728,8 @@ 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]; @@ -792,6 +856,14 @@ 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; } } @@ -823,6 +895,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 +969,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, @@ -965,6 +1051,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 +1072,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 +1133,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 +1203,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; @@ -1208,6 +1303,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 +1339,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 +1413,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, @@ -1423,6 +1537,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, @@ -1846,4 +1961,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/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index 0639fb14..fa3e6651 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -3168,6 +3168,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, @@ -3232,6 +3233,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 +3319,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)); 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..1c37809c 100644 --- a/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi +++ b/phoenix/crates/serving/xai-recsys-engine/xai_recsys_engine.pyi @@ -211,6 +211,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, @@ -311,6 +312,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-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/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 From b089ce64891f9c50fab73aa00dbe65acb82f198f Mon Sep 17 00:00:00 2001 From: CI agent Date: Mon, 17 Aug 2026 20:20:07 +0000 Subject: [PATCH 02/18] Open-source X Recommendation Algorithm --- .../service-lib/src/lib.rs | 2 + bdsm/README.md | 16 +- bdsm/runtime/score_results_sink_focal.py | 37 ++-- bdsm/runtime/sink_policy.yaml | 26 +-- bdsm/tests/test_sink_policy.py | 12 +- grox/config/config.py | 2 + grox/core/data_loaders/data_types.py | 57 ++++-- grox/core/data_loaders/kafka_loader.py | 30 +++- grox/core/data_loaders/post_mapper.py | 16 ++ grox/core/lm/post.py | 6 + grox/flows/ptos/classifier.py | 54 ++++++ grox/flows/ptos/plan_safety_ptos.py | 9 +- grox/flows/ptos/state.py | 1 + ...ety_ptos_adult_content_cross_validation.py | 162 ++++++++++++++++++ .../task_safety_ptos_safemodel_sex_nudity.py | 60 +++---- grox/flows/reply_spam/constants.py | 2 + grox/flows/reply_spam/generators.py | 17 +- grox/flows/reply_spam/task_filter.py | 4 +- grox/libs/kafka_cli/consumer.py | 6 + grox/libs/kafka_cli/multi_region_consumer.py | 6 + .../ads_brand_safety_vf_hydrator.rs | 20 ++- .../candidate_hydrators/gizmoduck_hydrator.rs | 16 ++ home-mixer/models/brand_safety.rs | 160 +++++++++++++++++ home-mixer/models/candidate.rs | 30 +++- home-mixer/params/config.rs | 2 +- home-mixer/params/param.rs | 7 + home-mixer/scorers/ranking_scorer.rs | 24 +++ .../client_events_kafka_side_effect.rs | 24 ++- .../phoenix_request_cache_side_effect.rs | 1 + home-mixer/util/urt/mod.rs | 1 + home-mixer/util/urt/new_tweets_pill.rs | 5 +- .../util/urt/reverse_chron_following/mod.rs | 1 + .../phoenixRankAllCandidateProcessor.strato | 79 ++++----- .../src/mm_embedding_client.rs | 18 +- 34 files changed, 751 insertions(+), 162 deletions(-) create mode 100644 grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py diff --git a/abuse-enforcement-service/service-lib/src/lib.rs b/abuse-enforcement-service/service-lib/src/lib.rs index 244ab59e..3055c182 100644 --- a/abuse-enforcement-service/service-lib/src/lib.rs +++ b/abuse-enforcement-service/service-lib/src/lib.rs @@ -1704,6 +1704,8 @@ pub async fn start_kafka_consumers( ) .with_wily_config(WilyConfig::default()) .with_ssl(sasl_ssl_config(&consumer_conn, &sasl_password)) + .with_enable_auto_offset_store(false) + .with_enable_auto_commit(false) .with_fetch_timeout_ms(10000); BatchConsumerConfig::new(kafka, SERVICE_NAME) diff --git a/bdsm/README.md b/bdsm/README.md index 8f072412..7b928949 100644 --- a/bdsm/README.md +++ b/bdsm/README.md @@ -106,15 +106,19 @@ 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`. + 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** (MIN_ACTIONS, dominant-head pick) - and the `enforcement_note` proto field. The template *strings* and the + 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 diff --git a/bdsm/runtime/score_results_sink_focal.py b/bdsm/runtime/score_results_sink_focal.py index 818d6103..520807d6 100644 --- a/bdsm/runtime/score_results_sink_focal.py +++ b/bdsm/runtime/score_results_sink_focal.py @@ -97,19 +97,19 @@ def _resolve_head_names(_model_version): } -def build_enforcement_note(head_scores_list, action_hist_list): +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 30 actions. 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). + 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 < 30: + if total < min_actions: return None bot_heads = [ @@ -220,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), @@ -358,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." @@ -1138,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: @@ -1165,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(user_head_scores, user_action_hist), + "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", @@ -1491,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( @@ -1637,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 c87e9be7..4ec9eaea 100644 --- a/bdsm/runtime/sink_policy.yaml +++ b/bdsm/runtime/sink_policy.yaml @@ -1,15 +1,18 @@ -version: "2026-08-13" +version: "2026-08-14" -# Public-release note: the per-head operating points below (the 2-D -# [tau, lambda] pairs, cusp_delta, and reply_spam_hard_suspend_tau) are -# REDACTED in the open-source export — they are replaced with an -# out-of-range 9.99 sentinel (these are probabilities in [0, 1], so a -# 9.99 threshold never fires and is plainly a stub, not a real value). -# The production values are configured internally and are not part of -# this release. Structure, head names, and the action map are real. +# 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] @@ -37,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 475346d9..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(): @@ -57,16 +58,17 @@ def _heads(**scores): 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) is None - assert m.build_enforcement_note(heads, None) is None - assert m.build_enforcement_note(heads, []) is None - assert m.build_enforcement_note(_heads(FollowBot=0.4, LegitimateUser=0.8), long) is None + 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) + note = m.build_enforcement_note(heads, long, gate) assert note == f"{m._REDACTED_TEMPLATE} [model: FollowBot=0.91]" diff --git a/grox/config/config.py b/grox/config/config.py index 87ad75da..2b9a5eab 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -89,6 +89,7 @@ class ModelName: 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" class NightOwlConfig(BaseModel): @@ -120,6 +121,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..a6a30571 100644 --- a/grox/core/data_loaders/data_types.py +++ b/grox/core/data_loaders/data_types.py @@ -367,30 +367,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\nTitle: {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 @@ -648,6 +665,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 @@ -744,6 +762,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 @@ -778,6 +802,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") diff --git a/grox/core/data_loaders/kafka_loader.py b/grox/core/data_loaders/kafka_loader.py index 60ed5276..7cd2e47c 100644 --- a/grox/core/data_loaders/kafka_loader.py +++ b/grox/core/data_loaders/kafka_loader.py @@ -13,6 +13,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 +28,7 @@ logger = logging.getLogger(__name__) MAX_WORKING_THREADS = 12 +_limiter = strategies.FixedWindowRateLimiter(storage.MemoryStorage()) def parse_native_encoding(data: bytes) -> int: @@ -55,6 +57,8 @@ def __init__(self, topic_name: str): self.consumer = KafkaConsumer(self.consumer_config) self.queue: asyncio.Queue[MessageQueuePayload] = asyncio.Queue() self._prefetcher_task: asyncio.Task | None = None + self._max_qps_per_partition = self.loader_config.max_qps_per_partition + self._limit_item: RateLimitItemPerSecond | None = None @staticmethod def _maybe_inject_scram_password(consumer_config) -> None: @@ -100,20 +104,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/post_mapper.py b/grox/core/data_loaders/post_mapper.py index ff432bf6..71111d5e 100644 --- a/grox/core/data_loaders/post_mapper.py +++ b/grox/core/data_loaders/post_mapper.py @@ -11,6 +11,7 @@ CardMetadataV2 as StratoCardMetadataV2, PollCardMetadata as StratoPollCardMetadata, GrokShareCardMetadata as StratoGrokShareCardMetadata, + GrokShareMetadata as StratoGrokShareMetadata, ArticleMetadata as StratoArticleMetadata, ListMetadata as StratoListMetadata, ChatGroupMetadata as StratoChatGroupMetadata, @@ -30,6 +31,7 @@ BroadcastMetadata, PollCard, GrokShareCard, + GrokShare, ArticleMetadata, ListMetadata, ChatGroupMetadata, @@ -117,6 +119,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( @@ -151,6 +159,7 @@ 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, @@ -303,6 +312,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 diff --git a/grox/core/lm/post.py b/grox/core/lm/post.py index 4df57298..e75b7cff 100644 --- a/grox/core/lm/post.py +++ b/grox/core/lm/post.py @@ -79,6 +79,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()) diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index 46805c11..b7dc5d21 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -82,12 +82,23 @@ def _fav_bucket(fav_count: int) -> str: 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, + recovery_timeout=600.0, + half_open_max_calls=5, + excluded_exceptions=(asyncio.CancelledError,), +) _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 +) class SafetyPtosCategoryClassifier: @@ -541,3 +552,46 @@ async def _sample(self, convo: Conversation, sample_for_gemma: bool = False) -> 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/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/state.py b/grox/flows/ptos/state.py index 9be56b7e..51a25a31 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 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..059aeb13 --- /dev/null +++ b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py @@ -0,0 +1,162 @@ +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 SAFETY_PTOS_DELUXE +from grox.flows.ptos.state import ( + SafetyPolicy, + SafetyPolicyCategory, + SafetyPolicyType, + SafetyPtosState, + SafetyPtosViolatedPolicy, +) +from monitor.metrics import Metrics +from strato_http.queries.safety_post_annotations_result import ( + StratoSafetyPostAnnotationsResultDirectMh, +) + +logger = logging.getLogger(__name__) + +_METRIC_PREFIX = "task.safety_ptos_adult_content_cross_validation" +_CROSS_VALIDATION_REASON = "Grok 4.5 Cross Validation disagreed" + + +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() + _result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() + + @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 == SAFETY_PTOS_DELUXE + 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 cls._post_is_already_flagged_nsfw(post): + Metrics.counter(f"{_METRIC_PREFIX}.skipped.count").add( + 1, attributes={"reason": "prior_nsfw"} + ) + return + await cls._cross_validate(ctx, post) + + @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(f"{_METRIC_PREFIX}.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 + ) + + @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"}) + 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": "hard" if is_hard else "soft"} + ) + logger.info( + f"Post {post.id}: grok 4.5 cross validation judged {judged.policyType.value}" + ) + if is_hard: + return + + 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=_CROSS_VALIDATION_REASON, + ) + ] + violations.append(adult_violations[0]) + for violation in adult_violations: + violation.safetyPolicy = SafetyPolicy( + policyType=SafetyPolicyType.AdultContentSexualSoft, + reason=_CROSS_VALIDATION_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_safemodel_sex_nudity.py b/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py index a7709b61..b5339f49 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,14 @@ 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 strato_http.queries.safety_post_annotations_result import ( + StratoSafetyPostAnnotationsResultDirectMh, +) logger = logging.getLogger(__name__) @@ -35,6 +34,8 @@ class TaskSafetyPtosSafemodelSexNudity(TaskWithPost): + _result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() + @classmethod async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: try: @@ -50,6 +51,20 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: ) logger.warning(f"Post {post.id}: safemodel failed: {e}") + @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(f"{_METRIC_PREFIX}.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 def _has_adult_content_suspicion(cls, ctx: TaskContext) -> bool: annotations = ctx.state(SafetyPtosState).annotations @@ -71,6 +86,12 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: ) return + if is_deluxe and await cls._post_is_already_flagged_nsfw(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 +144,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/reply_spam/constants.py b/grox/flows/reply_spam/constants.py index 5b60b6f6..0668e709 100644 --- a/grox/flows/reply_spam/constants.py +++ b/grox/flows/reply_spam/constants.py @@ -1,5 +1,7 @@ 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..a046e2f5 100644 --- a/grox/flows/reply_spam/generators.py +++ b/grox/flows/reply_spam/generators.py @@ -6,25 +6,32 @@ 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 - PLANS_TO_INJECT = { - PlanSpamComment.KEY, - PlanReplyRanking.KEY, - PlanCoordinatedSpam.KEY, - } + PLANS_TO_INJECT = {PlanSpamComment.KEY, PlanCoordinatedSpam.KEY} def _get_loader(self): return KafkaPostLoader(TOPIC_UNIFIED_POSTS) +@register +class ReplyRankingTaskGenerator(StreamTaskGenerator): + TASK_GENERATOR_TYPE = REPLY_RANKING + PLANS_TO_INJECT = {PlanReplyRanking.KEY} + + def _get_loader(self): + return KafkaPostLoader(TOPIC_UNIFIED_POSTS_V3) + + @register class ReplyRankingRecoveryTaskGenerator(StreamTaskGenerator): TASK_GENERATOR_TYPE = REPLY_RANKING_RECOVERY diff --git a/grox/flows/reply_spam/task_filter.py b/grox/flows/reply_spam/task_filter.py index 3bf626a8..5436ee55 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 = 30000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 40000 @override @classmethod @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 30000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 40000 @override @classmethod 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/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs b/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs index 2d271fa7..6dc75c55 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)); } } 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/models/brand_safety.rs b/home-mixer/models/brand_safety.rs index b225d8c9..92d7f94d 100644 --- a/home-mixer/models/brand_safety.rs +++ b/home-mixer/models/brand_safety.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::HashMap; use xai_x_thrift::tweet_safety_label::{SafetyLabel, SafetyLabelSource, SafetyLabelType}; @@ -62,6 +63,82 @@ pub fn compute_verdict( BrandSafetyVerdict::Safe } +pub(crate) const MEDIUM_RISK_LABELS_V2: &[SafetyLabelType] = &[ + SafetyLabelType::NSFW_HIGH_PRECISION, + SafetyLabelType::NSFW_HIGH_RECALL, + 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_V2, + SafetyLabelType::GROK_NSFA_EXPANDED_V2, + SafetyLabelType::NSFW_TEXT, +]; + +pub(crate) const LOW_RISK_LABELS_V2: &[SafetyLabelType] = &[ + SafetyLabelType::NSFA_LIMITED_INVENTORY, + SafetyLabelType::GROK_NSFA_LIMITED_V2, + SafetyLabelType::NSFA_HIGH_RECALL, +]; + +fn strip_v1_grok_written( + labels: &HashMap, +) -> Cow<'_, HashMap> { + const V1_DUAL_WRITTEN: &[SafetyLabelType] = &[ + SafetyLabelType::NSFA_HIGH_PRECISION, + SafetyLabelType::NSFA_LIMITED_INVENTORY, + ]; + const PROMPT_OWNED_RULES: &[i64] = &[1400, 1410, 1420, 1500, 1510, 1610, 1700]; + let is_v1_grox_written = |label_type: &SafetyLabelType, label: &SafetyLabel| { + V1_DUAL_WRITTEN.contains(label_type) + && botmaker_rule_id_from(label) + .is_some_and(|rule_id| PROMPT_OWNED_RULES.contains(&rule_id)) + }; + if !labels.iter().any(|(t, l)| is_v1_grox_written(t, l)) { + return Cow::Borrowed(labels); + } + Cow::Owned( + labels + .iter() + .filter(|&(t, l)| !is_v1_grox_written(t, l)) + .map(|(t, l)| (*t, l.clone())) + .collect(), + ) +} + +pub(crate) fn compute_verdict_v2( + labels: &HashMap, + tweet_id: u64, +) -> BrandSafetyVerdict { + let labels = strip_v1_grok_written(labels); + 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 @@ -213,6 +290,89 @@ mod tests { ); } + #[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) + .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) + .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..e19fa077 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)] @@ -62,6 +63,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)] @@ -72,12 +74,20 @@ pub struct PostCandidate { } #[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, } #[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] @@ -159,6 +169,13 @@ 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, }), reward_rerank_slot_prob: None, } @@ -182,6 +199,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 +231,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/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 8b0054e4..b553599e 100644 --- a/home-mixer/params/param.rs +++ b/home-mixer/params/param.rs @@ -917,6 +917,13 @@ param!( "rust_home_mixer_ads_blender", "partition_organic_low_risk" ); + +param!( + EnableAdsBrandSafetyVerdictV2, + bool, + "rust_home_mixer_ads_bs_v2_exp_enabled", + false +); param!( AdsTimeGapTSec, f64, diff --git a/home-mixer/scorers/ranking_scorer.rs b/home-mixer/scorers/ranking_scorer.rs index 6dc258ac..3139fb7a 100644 --- a/home-mixer/scorers/ranking_scorer.rs +++ b/home-mixer/scorers/ranking_scorer.rs @@ -658,17 +658,40 @@ impl RankingScorer { let mut contexts = vec![SlateContext::default(); candidates.len()]; let mut author_counts: FxHashMap = FxHashMap::default(); let mut last_author_rank: FxHashMap = FxHashMap::default(); + let mut sid_counts: [FxHashMap; 3] = Default::default(); + let mut last_sid_rank: [FxHashMap; 3] = Default::default(); for (rank, (idx, score)) in indexed.into_iter().enumerate() { let rank = rank as u32; 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); + + let mut sid_k = [0u32; 3]; + let mut sid_gap = [None; 3]; + let sids = candidates[idx].semantic_ids.as_deref().unwrap_or(&[]); + let sid_known = !sids.is_empty(); + let mut prefix = 0u64; + for (level, &code) in sids.iter().take(3).enumerate() { + prefix = (prefix << 20) | (code as u32 as u64 & 0xFFFFF); + sid_k[level] = sid_counts[level].get(&prefix).copied().unwrap_or(0); + sid_gap[level] = last_sid_rank[level].get(&prefix).map(|last| rank - last); + sid_counts[level].insert(prefix, sid_k[level] + 1); + last_sid_rank[level].insert(prefix, rank); + } + contexts[idx] = SlateContext { k, pool_rank: rank, pool_rank_gap: rank_gap, fatigue: 0.0, pre_diversity_score: score, + sid_known, + sid_k_l1: sid_k[0], + sid_k_l2: sid_k[1], + sid_k_l3: sid_k[2], + sid_gap_l1: sid_gap[0], + sid_gap_l2: sid_gap[1], + sid_gap_l3: sid_gap[2], }; author_counts.insert(author_id, k + 1); last_author_rank.insert(author_id, rank); @@ -1003,6 +1026,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), 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/phoenix_request_cache_side_effect.rs b/home-mixer/side_effects/phoenix_request_cache_side_effect.rs index 632affdb..a0563925 100644 --- a/home-mixer/side_effects/phoenix_request_cache_side_effect.rs +++ b/home-mixer/side_effects/phoenix_request_cache_side_effect.rs @@ -230,6 +230,7 @@ mod tests { pool_rank_gap: Some(3), fatigue: 1.5, pre_diversity_score: 0.42, + ..Default::default() }), ..candidate(1, 100) }; diff --git a/home-mixer/util/urt/mod.rs b/home-mixer/util/urt/mod.rs index 9b8e4924..102584b1 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 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/reverse_chron_following/mod.rs b/home-mixer/util/urt/reverse_chron_following/mod.rs index 9b0b257e..2a12825d 100644 --- a/home-mixer/util/urt/reverse_chron_following/mod.rs +++ b/home-mixer/util/urt/reverse_chron_following/mod.rs @@ -35,6 +35,7 @@ pub(crate) fn make_urt_timeline( cursor, language_code, country_code, + false, ) .into_iter() .collect(); diff --git a/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato b/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato index 7b92f8d2..6a0aaa87 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 = { @@ -340,39 +337,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 => () } 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..5221bf55 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 @@ -29,15 +29,6 @@ const BASE_PATH: &str = "/dev/shm/mm_embeddings"; const EMBEDDING_TTL: Duration = Duration::from_secs(2 * 24 * 3600); -fn embedding_ttl() -> Duration { - std::env::var("MM_EMBEDDING_TTL_SECS") - .ok() - .and_then(|v| v.trim().parse::().ok()) - .filter(|&secs| secs >= 1) - .map(Duration::from_secs) - .unwrap_or(EMBEDDING_TTL) -} - lazy_static! { static ref MM_EMBEDDING_SUCCESS_RATIO: Histogram = register_histogram!( "mm_embedding_success_ratio", @@ -226,7 +217,6 @@ pub fn get_mm_client( ) -> Result<(MmEmbeddingsClient, O2PreloadFuture)> { let is_writer = worker_id == 0; let shard_capacity = MAX_EMBEDDING_CACHE_SIZE / NUM_SHARDS; - let ttl = embedding_ttl(); let shards: Vec = if shared { let mut shards = Vec::with_capacity(NUM_SHARDS); @@ -244,10 +234,10 @@ pub fn get_mm_client( &data_path, shard_capacity, emb_dim, - ttl, + EMBEDDING_TTL, ) } else { - LmdbEmbeddingCache::open(&lmdb_path, &data_path, ttl) + LmdbEmbeddingCache::open(&lmdb_path, &data_path, EMBEDDING_TTL) }; CacheShard::Lmdb(Arc::new(cache.unwrap_or_else(|e| { panic!("LmdbEmbeddingCache shard {} failed: {}", shard_idx, e) @@ -267,7 +257,9 @@ pub fn get_mm_client( MAX_EMBEDDING_CACHE_SIZE ); (0..NUM_SHARDS) - .map(|_| CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, ttl)))) + .map(|_| { + CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, EMBEDDING_TTL))) + }) .collect() }; From 11a71f87d6a7fc4c1e8159dad8f3c5ff90a0f7ed Mon Sep 17 00:00:00 2001 From: CI agent Date: Tue, 18 Aug 2026 18:57:00 +0000 Subject: [PATCH 03/18] Open-source X Recommendation Algorithm --- ...ety_ptos_adult_content_cross_validation.py | 4 +- grox/flows/reply_spam/constants.py | 2 - grox/flows/reply_spam/generators.py | 17 +--- grox/flows/reply_spam/task_filter.py | 4 +- home-mixer/candidate_hydrators/mod.rs | 1 + .../vf_candidate_hydrator.rs | 2 +- .../vf_following_candidate_hydrator.rs | 95 ++++++++++++++++++ .../reverse_chron_posts_pipeline.rs | 7 +- home-mixer/models/brand_safety.rs | 98 +++++++++++++------ home-mixer/scorers/author_cold_start.rs | 17 ++-- home-mixer/scorers/vm_ranker.rs | 7 ++ .../postCreationEventProcessor.strato | 5 +- .../postCreationEventForwarder.strato | 16 +-- visibility-filtering/server_deps.rs | 3 + 14 files changed, 208 insertions(+), 70 deletions(-) create mode 100644 home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs 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 index 059aeb13..d42020bf 100644 --- a/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py +++ b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py @@ -127,9 +127,7 @@ async def _cross_validate(cls, ctx: TaskContext, post: Post) -> None: return is_hard = judged.policyType == SafetyPolicyType.AdultContentSexualHard - Metrics.counter(metric).add( - 1, attributes={"outcome": "hard" if is_hard else "soft"} - ) + 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}" ) diff --git a/grox/flows/reply_spam/constants.py b/grox/flows/reply_spam/constants.py index 0668e709..1a74a667 100644 --- a/grox/flows/reply_spam/constants.py +++ b/grox/flows/reply_spam/constants.py @@ -1,7 +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 a046e2f5..c41828e9 100644 --- a/grox/flows/reply_spam/generators.py +++ b/grox/flows/reply_spam/generators.py @@ -5,28 +5,21 @@ from grox.flows.reply_spam.plan_coordinated_spam import PlanCoordinatedSpam 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 - PLANS_TO_INJECT = {PlanSpamComment.KEY, PlanCoordinatedSpam.KEY} - - def _get_loader(self): - return KafkaPostLoader(TOPIC_UNIFIED_POSTS) - - @register class ReplyRankingTaskGenerator(StreamTaskGenerator): TASK_GENERATOR_TYPE = REPLY_RANKING - PLANS_TO_INJECT = {PlanReplyRanking.KEY} + PLANS_TO_INJECT = { + PlanReplyRanking.KEY, + PlanSpamComment.KEY, + PlanCoordinatedSpam.KEY, + } def _get_loader(self): return KafkaPostLoader(TOPIC_UNIFIED_POSTS_V3) diff --git a/grox/flows/reply_spam/task_filter.py b/grox/flows/reply_spam/task_filter.py index 5436ee55..ee3518aa 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 = 40000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 60000 @override @classmethod @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 40000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 60000 @override @classmethod diff --git a/home-mixer/candidate_hydrators/mod.rs b/home-mixer/candidate_hydrators/mod.rs index 6527acf1..a31d6539 100644 --- a/home-mixer/candidate_hydrators/mod.rs +++ b/home-mixer/candidate_hydrators/mod.rs @@ -17,3 +17,4 @@ 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/vf_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs index 24cca235..1d935657 100644 --- a/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs +++ b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs @@ -136,7 +136,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..c32e69dc --- /dev/null +++ b/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs @@ -0,0 +1,95 @@ +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 + }; + + 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/reverse_chron_posts_pipeline.rs b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs index 90cac898..534cc979 100644 --- a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs +++ b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs @@ -2,7 +2,7 @@ use crate::candidate_hydrators::ads_brand_safety_vf_hydrator::AdsBrandSafetyVfHy use crate::candidate_hydrators::conversation_gap_ancestor_hydrator::ConversationGapAncestorHydrator; use crate::candidate_hydrators::core_data_candidate_hydrator::CoreDataCandidateHydrator; 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}; @@ -138,7 +138,10 @@ impl ReverseChronPostsPipeline { ]; let post_selection_hydrators: Vec>> = vec![ - Box::new(VFCandidateHydrator::new(strato_vf_client, xai_vf_client).await), + Box::new(VFFollowingCandidateHydrator::new( + strato_vf_client, + xai_vf_client, + )), Box::new(AdsBrandSafetyVfHydrator { client: vf_safety_labels_client, }), diff --git a/home-mixer/models/brand_safety.rs b/home-mixer/models/brand_safety.rs index 92d7f94d..1ad616e7 100644 --- a/home-mixer/models/brand_safety.rs +++ b/home-mixer/models/brand_safety.rs @@ -1,4 +1,3 @@ -use std::borrow::Cow; use std::collections::HashMap; use xai_x_thrift::tweet_safety_label::{SafetyLabel, SafetyLabelSource, SafetyLabelType}; @@ -66,7 +65,6 @@ pub fn compute_verdict( pub(crate) const MEDIUM_RISK_LABELS_V2: &[SafetyLabelType] = &[ SafetyLabelType::NSFW_HIGH_PRECISION, SafetyLabelType::NSFW_HIGH_RECALL, - SafetyLabelType::NSFA_HIGH_PRECISION, SafetyLabelType::NSFA_KEYWORDS_HIGH_PRECISION, SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION, SafetyLabelType::NSFW_REPORTED_HEURISTICS, @@ -83,41 +81,24 @@ pub(crate) const MEDIUM_RISK_LABELS_V2: &[SafetyLabelType] = &[ ]; pub(crate) const LOW_RISK_LABELS_V2: &[SafetyLabelType] = &[ - SafetyLabelType::NSFA_LIMITED_INVENTORY, SafetyLabelType::GROK_NSFA_LIMITED_V2, SafetyLabelType::NSFA_HIGH_RECALL, ]; -fn strip_v1_grok_written( - labels: &HashMap, -) -> Cow<'_, HashMap> { - const V1_DUAL_WRITTEN: &[SafetyLabelType] = &[ - SafetyLabelType::NSFA_HIGH_PRECISION, - SafetyLabelType::NSFA_LIMITED_INVENTORY, - ]; - const PROMPT_OWNED_RULES: &[i64] = &[1400, 1410, 1420, 1500, 1510, 1610, 1700]; - let is_v1_grox_written = |label_type: &SafetyLabelType, label: &SafetyLabel| { - V1_DUAL_WRITTEN.contains(label_type) - && botmaker_rule_id_from(label) - .is_some_and(|rule_id| PROMPT_OWNED_RULES.contains(&rule_id)) - }; - if !labels.iter().any(|(t, l)| is_v1_grox_written(t, l)) { - return Cow::Borrowed(labels); - } - Cow::Owned( - labels - .iter() - .filter(|&(t, l)| !is_v1_grox_written(t, l)) - .map(|(t, l)| (*t, l.clone())) - .collect(), - ) -} +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 { - let labels = strip_v1_grok_written(labels); + if !V2_WRITTEN_LABELS.iter().any(|l| labels.contains_key(l)) { + return compute_verdict(labels, tweet_id); + } if MEDIUM_RISK_LABELS_V2.iter().any(|l| labels.contains_key(l)) { return BrandSafetyVerdict::MediumRisk; } @@ -290,6 +271,58 @@ 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 { @@ -354,7 +387,9 @@ mod tests { let expected_medium_v2: HashSet<_> = medium .iter() .copied() - .filter(|l| *l != SafetyLabelType::GROK_NSFA) + .filter(|l| { + *l != SafetyLabelType::GROK_NSFA && *l != SafetyLabelType::NSFA_HIGH_PRECISION + }) .chain([ SafetyLabelType::GROK_NSFA_V2, SafetyLabelType::GROK_NSFA_EXPANDED_V2, @@ -367,7 +402,10 @@ mod tests { let expected_low_v2: HashSet<_> = low .iter() .copied() - .filter(|l| *l != SafetyLabelType::GROK_NSFA_LIMITED) + .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); diff --git a/home-mixer/scorers/author_cold_start.rs b/home-mixer/scorers/author_cold_start.rs index e7652f85..2340e28c 100644 --- a/home-mixer/scorers/author_cold_start.rs +++ b/home-mixer/scorers/author_cold_start.rs @@ -102,12 +102,7 @@ fn count_tracked_ids( continue; } let source = c.served_type.map(|t| t as i32).unwrap_or(0); - if tracked.contains(&c.tweet_id) { - *counts.entry((c.tweet_id, source)).or_insert(0) += 1; - } - if tracked.contains(&c.author_id) { - *counts.entry((c.author_id, source)).or_insert(0) += 1; - } + *counts.entry((c.tweet_id, source)).or_insert(0) += 1; } counts } @@ -720,18 +715,22 @@ rust_home_mixer: ]; let counts = count_tracked_ids(&candidates, &tracked); assert_eq!( - counts.get(&(10, pb::ServedType::ForYouPhoenixRetrieval as i32)), + counts.get(&(1, pb::ServedType::ForYouPhoenixRetrieval as i32)), Some(&1) ); assert_eq!( - counts.get(&(10, pb::ServedType::ForYouPhoenixRetrievalMoe as i32)), + 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(&(30, 0)), None); + assert_eq!( + counts.get(&(10, pb::ServedType::ForYouPhoenixRetrieval as i32)), + None + ); + assert_eq!(counts.get(&(3, 0)), None); } #[test] diff --git a/home-mixer/scorers/vm_ranker.rs b/home-mixer/scorers/vm_ranker.rs index 624264b0..67207341 100644 --- a/home-mixer/scorers/vm_ranker.rs +++ b/home-mixer/scorers/vm_ranker.rs @@ -194,6 +194,13 @@ fn build_request(query: &ScoredPostsQuery, candidates: &[PostCandidate]) -> Rank k: s.k, pool_rank: s.pool_rank, pool_rank_gap: s.pool_rank_gap, + sid_known: s.sid_known, + sid_k1: s.sid_k_l1, + sid_k2: s.sid_k_l2, + sid_k3: s.sid_k_l3, + sid_gap1: s.sid_gap_l1, + sid_gap2: s.sid_gap_l2, + sid_gap3: s.sid_gap_l3, }), head_weights: scoring_weights .as_ref() 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/visibility-filtering/server_deps.rs b/visibility-filtering/server_deps.rs index 6e41332a..c47c86ae 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -241,12 +241,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() } } From aad7179773944e17eb8798bbbf0231d6cd6c1ffc Mon Sep 17 00:00:00 2001 From: CI agent Date: Wed, 19 Aug 2026 21:25:42 +0000 Subject: [PATCH 04/18] Open-source X Recommendation Algorithm --- grox/config/config.py | 1 + grox/flows/ptos/classifier.py | 73 +++- grox/flows/ptos/prior_nsfw.py | 34 ++ grox/flows/ptos/state.py | 1 + ...ety_ptos_adult_content_cross_validation.py | 21 +- grox/flows/ptos/task_safety_ptos_policy.py | 82 ++-- .../task_safety_ptos_safemodel_sex_nudity.py | 22 +- .../conversation_gap_ancestor_hydrator.rs | 28 +- home-mixer/candidate_hydrators/mod.rs | 1 + .../quoted_post_text_hydrator.rs | 106 +++++ .../phoenix_candidate_pipeline.rs | 4 +- .../reverse_chron_posts_pipeline.rs | 8 +- .../following_viewer_muted_keyword_filter.rs | 65 ++++ home-mixer/filters/mod.rs | 3 +- ...lter.rs => viewer_muted_keyword_filter.rs} | 30 +- home-mixer/models/candidate.rs | 2 + phoenix/Cargo.toml | 2 + phoenix/crates/common/xai-recsys/src/util.rs | 182 ++++++++- .../serving/xai-recsys-engine/Cargo.toml | 5 +- .../xai-recsys-engine/src/checkpoint_store.rs | 179 ++++++++- .../serving/xai-recsys-engine/src/python.rs | 54 +-- .../xai-recsys-engine/src/sid_client.rs | 45 +++ .../xai-recsys-engine/src/storage_util.rs | 9 +- .../serving/xai-recsys-engine/src/util.rs | 125 +++++- .../src/mm_embedding_client.rs | 18 +- .../xai-recsys-proto/proto/recsys.proto | 33 +- .../python/common/xai-proto/hatch_build.py | 8 + .../common/xai-proto/proto/recsys.proto | 33 +- phoenix/xrex/configs/data_feeds.py | 11 + phoenix/xrex/configs/xrecsys.py | 123 +++++- phoenix/xrex/configs/xrecsys_gen_recs.py | 8 +- phoenix/xrex/configs/xrecsys_sid_retrieval.py | 8 +- phoenix/xrex/configs/xrecsys_two_tower.py | 9 +- phoenix/xrex/cuda/async_emb/async_emb.py | 64 +++ .../xrex/cuda/async_emb/src/async_emb_api.cc | 283 ++++++++++++-- .../xrex/cuda/async_emb/src/async_emb_comm.cc | 15 +- .../cuda/async_emb/src/async_emb_comm.hpp | 6 +- .../cuda/async_emb/src/async_emb_kernel.cu | 21 +- .../cuda/async_emb/src/async_emb_kernel.hpp | 7 +- .../cutedsl/ranker_attention_varlen_fa4.py | 258 ++++++++---- .../xrex/cutedsl/ranker_fa4/block_sparsity.py | 37 ++ .../cutedsl/ranker_fa4/flash_bwd_sm100.py | 24 +- .../cutedsl/ranker_fa4/flash_fwd_sm100.py | 12 + phoenix/xrex/cutedsl/ranker_fa4/mask.py | 87 ++++- phoenix/xrex/data/grpc_recsys.py | 1 + phoenix/xrex/data/parquet_recsys.py | 2 + phoenix/xrex/data/recsys/constants.py | 16 + phoenix/xrex/data/recsys/feature_config.py | 51 ++- phoenix/xrex/data/recsys/recsys_batch.py | 48 ++- phoenix/xrex/data/retrieval_dataset.py | 22 +- phoenix/xrex/data/rust_kafka_recsys.py | 33 +- phoenix/xrex/data/rust_parquet_recsys.py | 1 + phoenix/xrex/data/streaming/kafkaloader.py | 142 ++++++- phoenix/xrex/eval/eval_utils.py | 3 +- phoenix/xrex/eval/metrics.py | 2 +- phoenix/xrex/inference/launch_inference.py | 30 +- phoenix/xrex/inference/model_runner.py | 367 ++++++++++-------- phoenix/xrex/inference/pinned_d2h.py | 13 + .../xrex/inference/sid_retrieval_runner.py | 19 +- phoenix/xrex/models/recsys_attention.py | 6 + phoenix/xrex/models/recsys_feature_prep.py | 85 +++- phoenix/xrex/models/recsys_model.py | 290 ++++++++++++-- phoenix/xrex/models/recsys_two_tower_model.py | 24 +- phoenix/xrex/optimizers/recsys/__init__.py | 1 + .../recsys/async_emb_gradient_update.py | 2 +- phoenix/xrex/optimizers/recsys/dense_optim.py | 47 +++ phoenix/xrex/optimizers/recsys/muon.py | 240 ++++++++++++ .../xrex/optimizers/recsys/rowwise_adagrad.py | 89 ++++- phoenix/xrex/settings.py | 2 +- phoenix/xrex/train/misc.py | 8 + phoenix/xrex/train/trainer.py | 49 ++- phoenix/xrex/train/trainer_recsys.py | 56 ++- phoenix/xrex/utils/log_timer.py | 3 +- phoenix/xrex/utils/metrics.py | 29 +- .../clients/gizmoduck_client.rs | 4 +- visibility-filtering/dark_traffic_setup.rs | 109 ++++++ .../hydration/tes_hydrator.rs | 8 + visibility-filtering/lib.rs | 1 + visibility-filtering/main.rs | 4 + 79 files changed, 3303 insertions(+), 651 deletions(-) create mode 100644 grox/flows/ptos/prior_nsfw.py create mode 100644 home-mixer/candidate_hydrators/quoted_post_text_hydrator.rs create mode 100644 home-mixer/filters/following_viewer_muted_keyword_filter.rs rename home-mixer/filters/{muted_keyword_filter.rs => viewer_muted_keyword_filter.rs} (92%) create mode 100644 phoenix/xrex/optimizers/recsys/dense_optim.py create mode 100644 phoenix/xrex/optimizers/recsys/muon.py create mode 100644 visibility-filtering/dark_traffic_setup.rs diff --git a/grox/config/config.py b/grox/config/config.py index 2b9a5eab..81e2065a 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -90,6 +90,7 @@ class ModelName: 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): diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index b7dc5d21..fade40f2 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -82,6 +82,14 @@ def _fav_bucket(fav_count: int) -> str: half_open_max_calls=5, excluded_exceptions=(asyncio.CancelledError,), ) +_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, @@ -99,6 +107,11 @@ def _fav_bucket(fav_count: int) -> str: _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 +) + +_GROK_4_6_INTERNAL_DIAL = 0.1 class SafetyPtosCategoryClassifier: @@ -189,6 +202,10 @@ 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())) + eapi_cfg_4_6 = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_6_INTERNAL) + self.eapi_4_6_internal = EapiSampler( + EapiModelConfig(**eapi_cfg_4_6.model_dump()) + ) def build_convo(self, post: Post) -> Conversation: post_creation_time = ( @@ -259,12 +276,18 @@ async def classify_policy(self, post: Post) -> SafetyPolicy: 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" + metric = "safety_ptos.child_safety_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 - ) + if random.random() < _GROK_4_6_INTERNAL_DIAL: + async with _eapi_4_6_internal_breaker.guard(): + raw = await self.eapi_4_6_internal.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) + else: + 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( @@ -331,6 +354,13 @@ def __init__( EapiModelConfig(**eapi_config_4_5_internal.model_dump()) ) + eapi_config_4_6_internal = grox_config.get_eapi_model( + ModelName.EAPI_GROK_4_6_INTERNAL + ) + self.eapi_4_6_internal = EapiSampler( + EapiModelConfig(**eapi_config_4_6_internal.model_dump()) + ) + @staticmethod def _get_policy_prompt(violation: SafetyPtosViolatedPolicy, post: Post) -> str: if violation.category == SafetyPolicyCategory.ViolentMedia: @@ -444,8 +474,12 @@ 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) + if random.random() < _GROK_4_6_INTERNAL_DIAL: + mode = "deluxe-4.6-internal" + result = await self._sample_4_6_internal(convo) + else: + mode = "deluxe-4.5-internal" + result = await self._sample_4_5_internal(convo) else: mode = "deluxe-4.3" result = await self._sample_4_3(convo) @@ -534,6 +568,31 @@ async def _sample_4_5_internal(self, convo: Conversation) -> str: convo.interleave(), conversation_id=convo.conversation_id ) + 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_6_fallback.count").add( + 1, attributes={"endpoint": breaker.name, "reason": "breaker_open"} + ) + logger.warning( + 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_6_fallback.count").add( + 1, attributes={"endpoint": breaker.name, "reason": "error"} + ) + logger.error( + 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 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 51a25a31..cdaa80ea 100644 --- a/grox/flows/ptos/state.py +++ b/grox/flows/ptos/state.py @@ -124,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_safety_ptos_adult_content_cross_validation.py b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py index d42020bf..afd886fe 100644 --- a/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py +++ b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py @@ -14,9 +14,7 @@ SafetyPtosViolatedPolicy, ) from monitor.metrics import Metrics -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__) @@ -40,7 +38,6 @@ def is_disagreement(self) -> bool: class TaskSafetyPtosAdultContentCrossValidation(TaskWithPost): _judge = SafetyPtosAdultContentCrossValidationJudge() - _result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() @classmethod async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: @@ -81,27 +78,13 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: ) if is_deluxe and outcome.is_disagreement: - if await cls._post_is_already_flagged_nsfw(post): + 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) - @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(f"{_METRIC_PREFIX}.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 - ) - @staticmethod def _compare_outcome( safemodel_positive: bool, ptos_positive: bool diff --git a/grox/flows/ptos/task_safety_ptos_policy.py b/grox/flows/ptos/task_safety_ptos_policy.py index d3343acd..83e9995e 100644 --- a/grox/flows/ptos/task_safety_ptos_policy.py +++ b/grox/flows/ptos/task_safety_ptos_policy.py @@ -19,9 +19,7 @@ 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__) @@ -64,24 +62,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 +78,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" @@ -139,13 +103,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 b5339f49..7a4417b1 100644 --- a/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py +++ b/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py @@ -14,9 +14,7 @@ from grox.core.tasks.task import Task, TaskWithPost, TaskResultCategory from monitor.metrics import Metrics from grox.flows.ptos.constants import SAFETY_PTOS_DELUXE -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__) @@ -34,8 +32,6 @@ class TaskSafetyPtosSafemodelSexNudity(TaskWithPost): - _result_direct_mh = StratoSafetyPostAnnotationsResultDirectMh() - @classmethod async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: try: @@ -51,20 +47,6 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: ) logger.warning(f"Post {post.id}: safemodel failed: {e}") - @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(f"{_METRIC_PREFIX}.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 def _has_adult_content_suspicion(cls, ctx: TaskContext) -> bool: annotations = ctx.state(SafetyPtosState).annotations @@ -86,7 +68,7 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: ) return - if is_deluxe and await cls._post_is_already_flagged_nsfw(post): + 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} ) 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/mod.rs b/home-mixer/candidate_hydrators/mod.rs index a31d6539..4f11398d 100644 --- a/home-mixer/candidate_hydrators/mod.rs +++ b/home-mixer/candidate_hydrators/mod.rs @@ -12,6 +12,7 @@ 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; 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_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index 3595411e..609ba9d6 100644 --- a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs +++ b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs @@ -43,7 +43,6 @@ 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; @@ -55,6 +54,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; @@ -354,7 +354,7 @@ 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 diff --git a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs index 534cc979..46a6da66 100644 --- a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs +++ b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs @@ -1,6 +1,7 @@ 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::quoted_post_text_hydrator::QuotedPostTextHydrator; use crate::candidate_hydrators::tweet_type_metrics_hydrator::TweetTypeMetricsHydrator; use crate::candidate_hydrators::vf_following_candidate_hydrator::VFFollowingCandidateHydrator; use crate::clients::night_owl_client::{MockNightOwlClient, NightOwlClient, ProdNightOwlClient}; @@ -8,6 +9,7 @@ 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::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; @@ -129,11 +131,15 @@ 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), ]; 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 9cb53f72..ce6231cc 100644 --- a/home-mixer/filters/mod.rs +++ b/home-mixer/filters/mod.rs @@ -8,10 +8,10 @@ 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; @@ -26,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 92% rename from home-mixer/filters/muted_keyword_filter.rs rename to home-mixer/filters/viewer_muted_keyword_filter.rs index 463e50eb..ec84c051 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![ @@ -281,7 +281,7 @@ mod tests { #[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/models/candidate.rs b/home-mixer/models/candidate.rs index e19fa077..e89dc326 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -33,6 +33,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, diff --git a/phoenix/Cargo.toml b/phoenix/Cargo.toml index 9cbc645e..d807269e 100644 --- a/phoenix/Cargo.toml +++ b/phoenix/Cargo.toml @@ -57,9 +57,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" diff --git a/phoenix/crates/common/xai-recsys/src/util.rs b/phoenix/crates/common/xai-recsys/src/util.rs index 4ace24eb..3e446130 100644 --- a/phoenix/crates/common/xai-recsys/src/util.rs +++ b/phoenix/crates/common/xai-recsys/src/util.rs @@ -26,7 +26,7 @@ lazy_static! { .unwrap(); } -fn record_sid_coverage(sequence: &str, present: u64, count: u64) { +pub fn record_sid_coverage(sequence: &str, present: u64, count: u64) { SID_COVERAGE_TOTAL .with_label_values(&[sequence, "present"]) .inc_by(present); @@ -35,6 +35,18 @@ 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, @@ -397,15 +409,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 { @@ -732,6 +741,8 @@ impl InputBuffer { 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 { @@ -773,6 +784,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; @@ -868,6 +886,13 @@ impl InputBuffer { } } + 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); @@ -994,7 +1019,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, } @@ -1236,19 +1261,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(), + ); } } } @@ -1780,6 +1803,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; diff --git a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml index 6e0cecf4..99343fc9 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"] } @@ -69,6 +69,9 @@ xai-recsys-mm-server = { workspace = true } 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/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/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index fa3e6651..650a11ec 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -1171,36 +1171,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, @@ -2751,6 +2721,7 @@ struct RecsysRetrievalPredictorImpl { reload_directive: Arc>>, enqueue_timeout_ms: u64, mm_embeddings_client: Option, + #[allow(dead_code)] sid_client: Option>, admission: Arc, prefetch_mm_query_config: Option>, @@ -2870,7 +2841,6 @@ impl RecsysRetrievalPredictorImpl { &self.model_config, &cancel_token, self.enqueue_timeout_ms, - &self.sid_client, request.client_context, request.user_context, deadline, @@ -2975,7 +2945,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 +2959,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 +2998,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()); diff --git a/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs index 50b92704..be094e00 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs @@ -51,6 +51,36 @@ pub struct SemanticIdClient { sid_num_levels: usize, } +pub fn fill_semantic_ids( + sids: &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; + } + } + } +} + impl SemanticIdClient { pub fn new(endpoint: &str, sid_num_levels: usize) -> Self { let clients = (0..NUM_CHANNELS) @@ -142,3 +172,18 @@ impl PySemanticIdClient { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn semantic_ids_are_shifted_and_missing_values_remain_padding() { + let sids = HashMap::from([(11, vec![0, 7, -1]), (22, vec![3, 4, 5])]); + let mut dst = vec![0u16; 3 * 3]; + + fill_semantic_ids(&sids, &[11, 0, 22], &mut dst, 3); + + assert_eq!(dst, vec![1, 8, 0, 0, 0, 0, 4, 5, 6]); + } +} 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-mm-server/src/mm_embedding_client.rs b/phoenix/crates/serving/xai-recsys-mm-server/src/mm_embedding_client.rs index 5221bf55..c5313efc 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 @@ -29,6 +29,15 @@ const BASE_PATH: &str = "/dev/shm/mm_embeddings"; const EMBEDDING_TTL: Duration = Duration::from_secs(2 * 24 * 3600); +fn embedding_ttl() -> Duration { + std::env::var("MM_EMBEDDING_TTL_SECS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|&secs| secs >= 1) + .map(Duration::from_secs) + .unwrap_or(EMBEDDING_TTL) +} + lazy_static! { static ref MM_EMBEDDING_SUCCESS_RATIO: Histogram = register_histogram!( "mm_embedding_success_ratio", @@ -217,6 +226,7 @@ pub fn get_mm_client( ) -> Result<(MmEmbeddingsClient, O2PreloadFuture)> { let is_writer = worker_id == 0; let shard_capacity = MAX_EMBEDDING_CACHE_SIZE / NUM_SHARDS; + let ttl = embedding_ttl(); let shards: Vec = if shared { let mut shards = Vec::with_capacity(NUM_SHARDS); @@ -234,10 +244,10 @@ pub fn get_mm_client( &data_path, shard_capacity, emb_dim, - EMBEDDING_TTL, + ttl, ) } else { - LmdbEmbeddingCache::open(&lmdb_path, &data_path, EMBEDDING_TTL) + LmdbEmbeddingCache::open(&lmdb_path, &data_path, ttl) }; CacheShard::Lmdb(Arc::new(cache.unwrap_or_else(|e| { panic!("LmdbEmbeddingCache shard {} failed: {}", shard_idx, e) @@ -257,9 +267,7 @@ pub fn get_mm_client( MAX_EMBEDDING_CACHE_SIZE ); (0..NUM_SHARDS) - .map(|_| { - CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, EMBEDDING_TTL))) - }) + .map(|_| CacheShard::InProcess(Arc::new(EmbeddingCache::new(shard_capacity, ttl)))) .collect() }; diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index ca21eecf..efdb39f8 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -224,6 +224,14 @@ message NextActionDistribution { repeated float continuousActionsValues = 6; map indexToContinuousValues = 7; + + RewardOutputs rewardOutputs = 8; +} + +message RewardOutputs { + repeated float values = 1; + + repeated float policyLogits = 2; } enum ActionName { @@ -332,6 +340,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,16 +376,16 @@ 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_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; @@ -429,7 +439,7 @@ enum ActionName { 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; + P_OPEN_LINK_P90 = 253; PLACE_HOLDER_254 = 254; PLACE_HOLDER_255 = 255; } @@ -1192,6 +1202,7 @@ message ScoreInfo { optional double weightedScore = 2; optional double finalScore = 3; SlateContext slateContext = 4; + optional double rewardRerankSlotProb = 5; } message SlateContext { @@ -1200,6 +1211,13 @@ 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; } message ActionInfo { @@ -1635,6 +1653,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/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..efdb39f8 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -224,6 +224,14 @@ message NextActionDistribution { repeated float continuousActionsValues = 6; map indexToContinuousValues = 7; + + RewardOutputs rewardOutputs = 8; +} + +message RewardOutputs { + repeated float values = 1; + + repeated float policyLogits = 2; } enum ActionName { @@ -332,6 +340,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,16 +376,16 @@ 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_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; @@ -429,7 +439,7 @@ enum ActionName { 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; + P_OPEN_LINK_P90 = 253; PLACE_HOLDER_254 = 254; PLACE_HOLDER_255 = 255; } @@ -1192,6 +1202,7 @@ message ScoreInfo { optional double weightedScore = 2; optional double finalScore = 3; SlateContext slateContext = 4; + optional double rewardRerankSlotProb = 5; } message SlateContext { @@ -1200,6 +1211,13 @@ 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; } message ActionInfo { @@ -1635,6 +1653,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/xrex/configs/data_feeds.py b/phoenix/xrex/configs/data_feeds.py index d25c81d3..36e08cff 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,7 @@ 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), ) @@ -144,6 +147,7 @@ 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), ) @@ -173,6 +177,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 +197,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 +217,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"], diff --git a/phoenix/xrex/configs/xrecsys.py b/phoenix/xrex/configs/xrecsys.py index 787061dd..018d7e8c 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, + ), + ), } @@ -393,6 +435,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 +527,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 +612,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, @@ -554,6 +638,9 @@ 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 ), + condition_search_relevance_on_prompt=mparams.get( + "condition_search_relevance_on_prompt", False + ), metric_group=mparams.get("metric_group", "default"), continuous_metrics_mae_mean=mparams.get("continuous_metrics_mae_mean", False), emb_table_width=mparams["emb_table_width"], @@ -705,23 +792,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..ea4875b0 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 @@ -717,7 +716,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 +732,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_varlen_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py index c0dc49e1..d850e1b1 100644 --- a/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py +++ b/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py @@ -33,12 +33,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 +61,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 +68,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 +189,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 +198,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 +207,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 +227,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 +285,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 +323,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 +382,8 @@ def launch_fwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mO, mLSE, softmax_scale: cutlass.Float32, @@ -295,6 +400,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 +462,8 @@ def launch_bwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mdQa, mdKa, mdVa, @@ -372,6 +481,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 +516,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 +537,8 @@ def launch_bwd( mFullIdx, mDiagCnt, mDiagIdx, + mValidUpper, + mValidLower, mdQa, mdK, mdV, @@ -443,6 +556,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 +591,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 +662,22 @@ 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) + valid_bounds = bs_args[12:] + out, lse = c["fwd_call"](q, k, v, *fbs, *valid_bounds) out = checkpoint_name(out, "attn_outputs") lse = checkpoint_name(lse, "attn_outputs") 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 +698,7 @@ def _attention_bwd(res, g): lse_log2, dpsum, *bbs, + *valid_bounds, dq_accum_init, dk_accum_init, dv_accum_init, @@ -596,6 +715,7 @@ def _attention_bwd(res, g): lse_log2, dpsum, *bbs, + *valid_bounds, dq_accum_init, ) (dq,) = c["post_dq_call"](dq_accum) @@ -605,4 +725,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..debf9b90 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) diff --git a/phoenix/xrex/data/parquet_recsys.py b/phoenix/xrex/data/parquet_recsys.py index d35abe31..28150d59 100644 --- a/phoenix/xrex/data/parquet_recsys.py +++ b/phoenix/xrex/data/parquet_recsys.py @@ -868,6 +868,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 +1154,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: diff --git a/phoenix/xrex/data/recsys/constants.py b/phoenix/xrex/data/recsys/constants.py index 7d3e6868..17c7359c 100644 --- a/phoenix/xrex/data/recsys/constants.py +++ b/phoenix/xrex/data/recsys/constants.py @@ -179,6 +179,21 @@ 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", + ], +} + ads_conversion_engagement_to_action_types = { **primary_engagement_to_action_types, "IsPurchaseConversion": [ @@ -243,6 +258,7 @@ 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, "none": {}, diff --git a/phoenix/xrex/data/recsys/feature_config.py b/phoenix/xrex/data/recsys/feature_config.py index dd4d7bb3..0a287a2a 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,25 @@ class UserInt64Feature(enum.IntEnum): ] +OPTIONAL_COLUMNS: list[str] = [ + "searchQueryEmbeddingSeq", + "lineItemObjectiveSeq", + "safetyLabelMaskSeq", + "semanticIdSeq", + "sampleWeight", + "firstDpaProductKey", + "authorFollowerCountSeq", + "inReplyToPostIdSeq", +] + + 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 +170,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..9e180030 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) @@ -347,6 +357,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 @@ -832,6 +843,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, diff --git a/phoenix/xrex/data/retrieval_dataset.py b/phoenix/xrex/data/retrieval_dataset.py index 6a4af404..37335d5b 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,7 +47,7 @@ def _bridge_ads_s3_env() -> None: os.environ[generic] = val -_bridge_ads_s3_env() +_bridge_o2_env() def _idx(sub: str) -> str: @@ -96,11 +96,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()) 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..d2bec88d 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, ) @@ -825,6 +934,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/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..00f41e9e 100644 --- a/phoenix/xrex/inference/launch_inference.py +++ b/phoenix/xrex/inference/launch_inference.py @@ -463,7 +463,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", @@ -647,9 +656,9 @@ def run( "--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", @@ -791,9 +800,8 @@ def run( 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." + "Optional leftover SID lookup endpoint. History SIDs are parsed " + "from the request; this is not required for use_post_sid=True." ), ) parser.add_argument( @@ -845,5 +853,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..1255922e 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" @@ -638,6 +641,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 +863,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 +1499,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 +1535,7 @@ 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 _coordinator_loop(self) -> None: logger.info("[hotswap] Coordinator thread started (subprocess loader).") @@ -1535,97 +1560,103 @@ 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 + self._reload_requested.clear() + logger.info("[hotswap] Coordinator: reload triggered, sending to subprocess.") + bg_start = time.time() - 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"}) - 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() + 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() + + finally: + self._hotswap_cycle_inflight.clear() logger.info("[hotswap] Coordinator thread exiting.") @@ -1676,68 +1707,76 @@ 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() + 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 +3104,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 +3158,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", @@ -3229,11 +3275,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 +3311,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 +4081,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, @@ -4661,7 +4704,7 @@ def create_server( 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 + sid_num_levels = self.model_config.user_tower_config.sid_num_levels if _use_post_sid else 0 if _use_post_sid and self.sid_endpoint and sid_num_levels > 0: sid_client = xai_recsys_engine.PySemanticIdClient( self.sid_endpoint, @@ -4672,8 +4715,11 @@ def create_server( 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") + elif _use_post_sid: + logger.info( + "Parsing history SIDs from the request (sid_num_levels=%d); no sid_endpoint", + sid_num_levels, + ) return xai_recsys_engine.RecsysRetrievalPredictorServer( self.grpc_port, @@ -4720,10 +4766,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..c20149dd 100644 --- a/phoenix/xrex/inference/sid_retrieval_runner.py +++ b/phoenix/xrex/inference/sid_retrieval_runner.py @@ -378,20 +378,21 @@ def create_server( 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_num_levels = self.model_config.sid_num_levels if self.model_config.use_post_sid else 0 + if self.sid_endpoint and sid_num_levels > 0: sid_client = xai_recsys_engine.PySemanticIdClient( self.sid_endpoint, - self.model_config.sid_num_levels, + 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, + ) + elif self.model_config.use_post_sid: + logger.info( + "Parsing history SIDs from the request (sid_num_levels=%d); no sid_endpoint", + sid_num_levels, ) return xai_recsys_engine.RecsysRetrievalPredictorServer( @@ -442,5 +443,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..77c2c06f 100644 --- a/phoenix/xrex/models/recsys_attention.py +++ b/phoenix/xrex/models/recsys_attention.py @@ -265,6 +265,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 +292,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 +314,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..81b4a928 100644 --- a/phoenix/xrex/models/recsys_model.py +++ b/phoenix/xrex/models/recsys_model.py @@ -22,10 +22,17 @@ CLICK_ACTION_INDEX, CLICK_CONDITIONED_ACTION_INDICES, NEGATIVE_FEEDBACK_HEAD_INDICES, + SEARCH_RELEVANCE_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 @@ -90,6 +97,11 @@ ) +_DPA_PRODUCT_KEY_SLOTS: list[int] = [ + Int64Feature.firstDpaProductKey.value, + Int64Feature.firstDpaProductKeyHash2.value, +] + POST_AGE_MAX_MINUTES = 4800 @@ -498,6 +510,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 @@ -537,8 +550,16 @@ class RecsysAggregatedModelConfig(Config): condition_conversion_on_click: bool = False + 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() @@ -1117,6 +1138,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 +1151,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 +1223,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 +1243,7 @@ def pad_to_next_128_multiple( padded_client_app_id, padded_line_item_objective, padded_safety_label_mask, + padded_dpa_product_key, ) @@ -1229,8 +1261,10 @@ 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, *, condition_conversion_on_click: bool = False, + condition_search_relevance_on_prompt: bool = False, enable_platform_metrics: bool = False, ) -> dict[str, jax.Array]: promoted_mask = mask * (promoted_ids != 0) if promoted_ids is not None else jnp.zeros_like(mask) @@ -1307,11 +1341,26 @@ 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 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 @@ -1381,7 +1430,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 +1470,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 +1554,7 @@ 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, ) -> dict[str, jax.Array]: return build_metric_masks( mask, @@ -1506,7 +1566,9 @@ def _build_metric_masks( new_user_mask, line_item_objective, no_history_mask, + dpa_product_key, 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, ) @@ -1522,6 +1584,7 @@ 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, stats: dict | None = None, rce_ema: dict[str, jax.Array] | None = None, rce_alpha: jax.Array | None = None, @@ -1542,6 +1605,7 @@ def compute_recsys_metrics( new_user_mask, line_item_objective, no_history_mask, + dpa_product_key, ) return self._compute_metrics_after_masks( @@ -1612,29 +1676,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 +1827,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 +2069,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 +2242,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 +2288,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 +2360,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 +2418,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 +2505,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: @@ -2442,6 +2660,9 @@ def _embed_post_sid(seq_name: str) -> jnp.ndarray | None: fprop_dtype=DTYPE_BY_NAME[self.config.fprop_dtype], sid_post_embeddings=_sid_post_emb_c, ) + candidate_embeddings = self._maybe_add_dpa_input_embedding( + candidate_embeddings, recsys_features_batch + ) if _config.user_features.has_user_features: feature_parts = build_user_feature_parts( @@ -2596,6 +2817,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 ) @@ -2708,6 +2935,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 +2948,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, :] @@ -2831,6 +3060,15 @@ def loss( conv_zero_mask = no_click[:, :, None] * conv_head_mask loss_mask = loss_mask * (1 - conv_zero_mask) + 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) + safety_stats = safety_filter_stats( candidate_safety_mask, target_padding_mask, @@ -2877,6 +3115,7 @@ 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, stats=stats, rce_ema=rce_ema, rce_alpha=rce_alpha, @@ -2911,6 +3150,7 @@ 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, ) for loss_config in self.config.continuous_action_losses: diff --git a/phoenix/xrex/models/recsys_two_tower_model.py b/phoenix/xrex/models/recsys_two_tower_model.py index e142e9c2..fb14083c 100644 --- a/phoenix/xrex/models/recsys_two_tower_model.py +++ b/phoenix/xrex/models/recsys_two_tower_model.py @@ -1245,19 +1245,17 @@ 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 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/misc.py b/phoenix/xrex/train/misc.py index fba69fb8..08aaac4e 100644 --- a/phoenix/xrex/train/misc.py +++ b/phoenix/xrex/train/misc.py @@ -20,6 +20,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 diff --git a/phoenix/xrex/train/trainer.py b/phoenix/xrex/train/trainer.py index fc9b7381..ed7b4058 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -1149,6 +1149,10 @@ 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 maybe_load_checkpoint( self, ctx: TrainerContext, tag: str | None = None ) -> tuple[bool, int, int]: @@ -1176,8 +1180,7 @@ def maybe_load_checkpoint( ) and ctx.checkpoint.is_manual_load() 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 @@ -1245,8 +1248,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 +1448,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..78c942cc 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) @@ -1432,7 +1439,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, @@ -1522,6 +1529,7 @@ def _update(updated, original): "emb_grad_norm": emb_grad_norm, "emb_valid_step": emb_valid_step, } + metrics.update(emb_optim_metrics) if self.track_norm_metrics: metrics.update(norm_metrics(new_params, new_opt_state, gradients, updates)) @@ -1663,6 +1671,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)} @@ -1754,7 +1764,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 +2152,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"): @@ -2391,6 +2429,12 @@ 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 maybe_load_checkpoint(self, ctx: TrainerContext, tag=None): assert isinstance( self.model_config, (RecsysAggregatedModelConfig, RecsysTwoTowerModelConfig) @@ -2947,12 +2991,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 +3012,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) 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/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/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/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs new file mode 100644 index 00000000..5cb7486f --- /dev/null +++ b/visibility-filtering/dark_traffic_setup.rs @@ -0,0 +1,109 @@ +use std::sync::Arc; +use tonic::async_trait; +use tower::util::Either; +use tracing::info; + +use xai_dark_traffic::{DarkTrafficLayer, ReloadableDarkTrafficConfigBuilder}; +use xai_x_rpc::dynamic_channel_manager::{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"; +const SHADOW_WORKLOAD: &str = "xai-vf-shadow"; + +pub type DarkLayer = Either; + +struct StaticShadowDiscovery; + +#[async_trait] +impl EndpointDiscovery for StaticShadowDiscovery { + async fn discover(&self) -> anyhow::Result> { + Ok(vec![EndpointInfo { + name: SHADOW_WORKLOAD.to_string(), + xds_dest: format!("{SHADOW_WORKLOAD}.prod.visibility:grpc"), + }]) + } +} + +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 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_enable(ordinal, max_ordinal) { + info!( + ?ordinal, + ?max_ordinal, + "dark_traffic: disabled (ordinal >= max)" + ); + return Either::Right(tower::layer::util::Identity::new()); + } + + let dc = std::env::var("DATACENTER").unwrap_or_else(|_| "atla".to_string()); + let domain = format!("visibility.visibility-filtering-service.prod.{dc}.s2s.twttr.net"); + info!(domain, "dark_traffic: enabled"); + + let factory = XdsChannelFactory::new( + TlsMode::mtls_from_env() + .expect("S2S TLS config required") + .with_domain_override(&domain), + ); + + let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(StaticShadowDiscovery)); + + let config = ReloadableDarkTrafficConfigBuilder::new(CONFIG_PATH) + .forwarders({ + let ch = Arc::clone(&channels); + move || ch.channels() + }) + .build(); + + Either::Left(DarkTrafficLayer::new(config)) +} + +fn should_enable(ordinal: Option, max_ordinal: Option) -> bool { + let max = max_ordinal.unwrap_or(1); + let ord = ordinal.unwrap_or(u32::MAX); + ord < max +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_only_pod0() { + assert!(should_enable(Some(0), None)); + assert!(!should_enable(Some(1), None)); + assert!(!should_enable(Some(99), None)); + } + + #[test] + fn no_ordinal_disables() { + assert!(!should_enable(None, None)); + assert!(!should_enable(None, Some(3))); + } + + #[test] + fn max_ordinal_threshold() { + assert!(should_enable(Some(0), Some(3))); + assert!(should_enable(Some(2), Some(3))); + assert!(!should_enable(Some(3), Some(3))); + assert!(!should_enable(Some(4), Some(3))); + } + + #[test] + fn max_ordinal_zero_disables_all() { + assert!(!should_enable(Some(0), Some(0))); + } +} diff --git a/visibility-filtering/hydration/tes_hydrator.rs b/visibility-filtering/hydration/tes_hydrator.rs index 220e2171..35ff5a93 100644 --- a/visibility-filtering/hydration/tes_hydrator.rs +++ b/visibility-filtering/hydration/tes_hydrator.rs @@ -666,6 +666,14 @@ mod tests { unreachable!() } + async fn get_trusted_friends_controls( + &self, + _tweet_ids: Vec, + ) -> HashMap>> + { + unreachable!() + } + async fn get_grok_post_ids( &self, _tweet_ids: Vec, diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index e83fc76d..9809b07b 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -1,5 +1,6 @@ pub 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; diff --git a/visibility-filtering/main.rs b/visibility-filtering/main.rs index 0f47c471..f404a789 100644 --- a/visibility-filtering/main.rs +++ b/visibility-filtering/main.rs @@ -1,6 +1,8 @@ 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; @@ -30,6 +32,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 From d0cef2f943084ee0d4310378031c9c2c37d67f12 Mon Sep 17 00:00:00 2001 From: CI agent Date: Thu, 20 Aug 2026 20:22:58 +0000 Subject: [PATCH 05/18] Open-source X Recommendation Algorithm --- grox/core/data_loaders/data_types.py | 4 + grox/core/data_loaders/post_mapper.py | 3 + grox/flows/ptos/classifier.py | 2 +- grox/libs/grok_sampler/llm.py | 40 +- .../ai_trend_feedback_context_hydrator.rs | 157 ++++++++ home-mixer/candidate_hydrators/mod.rs | 1 + .../phoenix_candidate_pipeline.rs | 4 + home-mixer/models/candidate.rs | 23 ++ home-mixer/params/param.rs | 13 + home-mixer/scored_posts_server.rs | 2 + home-mixer/scorers/phoenix_scorer.rs | 4 + home-mixer/scorers/ranking_scorer.rs | 32 +- home-mixer/util/urt/post_marshaller.rs | 36 +- .../xai-recsys-engine/src/emb_table.rs | 358 +++++++++--------- 14 files changed, 479 insertions(+), 200 deletions(-) create mode 100644 home-mixer/candidate_hydrators/ai_trend_feedback_context_hydrator.rs diff --git a/grox/core/data_loaders/data_types.py b/grox/core/data_loaders/data_types.py index a6a30571..b99e3045 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, ) diff --git a/grox/core/data_loaders/post_mapper.py b/grox/core/data_loaders/post_mapper.py index 71111d5e..e9b05973 100644 --- a/grox/core/data_loaders/post_mapper.py +++ b/grox/core/data_loaders/post_mapper.py @@ -347,6 +347,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/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index fade40f2..b353e40e 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -111,7 +111,7 @@ def _fav_bucket(fav_count: int) -> str: ModelName.EAPI_GROK_4_6_INTERNAL, _EAPI_4_6_INTERNAL_BREAKER_CONFIG ) -_GROK_4_6_INTERNAL_DIAL = 0.1 +_GROK_4_6_INTERNAL_DIAL = 0.3 class SafetyPtosCategoryClassifier: 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/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/mod.rs b/home-mixer/candidate_hydrators/mod.rs index 4f11398d..9796a9b6 100644 --- a/home-mixer/candidate_hydrators/mod.rs +++ b/home-mixer/candidate_hydrators/mod.rs @@ -1,4 +1,5 @@ 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; diff --git a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index 609ba9d6..c1c0c9cc 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; @@ -423,6 +424,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/models/candidate.rs b/home-mixer/models/candidate.rs index e89dc326..f6510138 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -23,6 +23,8 @@ pub struct PostCandidate { pub score: Option, pub slate_context: Option, #[serde(default)] + pub served_slate_context: Option, + #[serde(default)] pub mpn_parts: Option, #[serde( serialize_with = "serialize_served_type", @@ -73,6 +75,8 @@ 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)] @@ -92,6 +96,25 @@ pub struct SlateContext { pub sid_gap_l3: Option, } +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, + } + } +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MpnParts { pub pos: f64, diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index b553599e..845c06c7 100644 --- a/home-mixer/params/param.rs +++ b/home-mixer/params/param.rs @@ -243,6 +243,12 @@ param!( "rust_home_mixer_log_slate_context", false ); +param!( + UseServedSlateContext, + bool, + "rust_home_mixer_use_served_slate_context", + false +); param!( OonWeightFactor, f64, @@ -1020,6 +1026,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..0c417b65 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() diff --git a/home-mixer/scorers/phoenix_scorer.rs b/home-mixer/scorers/phoenix_scorer.rs index 4ff7b1c8..5cf292d7 100644 --- a/home-mixer/scorers/phoenix_scorer.rs +++ b/home-mixer/scorers/phoenix_scorer.rs @@ -107,6 +107,9 @@ impl Scorer for PhoenixScorer { .iter() .map(|c| PostCandidate { phoenix_scores: predictions.candidate_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, ..Default::default() @@ -117,6 +120,7 @@ impl Scorer for PhoenixScorer { fn update(&self, candidate: &mut PostCandidate, scored: PostCandidate) { candidate.phoenix_scores = scored.phoenix_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; } diff --git a/home-mixer/scorers/ranking_scorer.rs b/home-mixer/scorers/ranking_scorer.rs index 3139fb7a..962ab7ae 100644 --- a/home-mixer/scorers/ranking_scorer.rs +++ b/home-mixer/scorers/ranking_scorer.rs @@ -700,6 +700,16 @@ impl RankingScorer { contexts } + fn served_slate_contexts( + query: &ScoredPostsQuery, + candidates: &[PostCandidate], + ) -> Option> { + if !query.params.get(UseServedSlateContext) { + return None; + } + candidates.iter().map(|c| c.served_slate_context).collect() + } + fn stored_slate_contexts(candidates: &[PostCandidate]) -> Option> { candidates.iter().map(|c| c.slate_context).collect() } @@ -806,11 +816,12 @@ impl Scorer for RankingScorer { }; 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> = + match Self::served_slate_contexts(query, candidates) { + Some(served) => Some(served), + None if query.has_cached_posts => Self::stored_slate_contexts(candidates), + None => Some(Self::compute_slate_contexts(candidates, &weighted_scores)), + }; let diversity_multipliers: Vec = if enable_author_diversity { let recomputed_contexts; @@ -875,11 +886,12 @@ 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 persisted_contexts: Option> = + match Self::served_slate_contexts(query, candidates) { + Some(served) => Some(served), + None if query.has_cached_posts => Self::stored_slate_contexts(candidates), + None => Some(Self::compute_slate_contexts(candidates, &adjusted_scores)), + }; let diversity_adjusted = if enable_author_diversity { let recomputed_contexts; 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/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs index 64a73542..0673dbbe 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs @@ -1161,63 +1161,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 +1230,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 +1247,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::*; From 28e414f535e4b5a50ca12ee87674e7649e50c7ad Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 21 Aug 2026 19:10:38 +0000 Subject: [PATCH 06/18] Open-source X Recommendation Algorithm --- grox/config/config.py | 1 - grox/flows/ptos/classifier.py | 69 +-------- grox/flows/reply_spam/task_filter.py | 4 +- phoenix/Cargo.lock | 70 +++++++++ phoenix/Cargo.toml | 2 + phoenix/README.md | 8 +- phoenix/TRAINING.md | 26 ++-- .../serving/xai-recsys-engine/Cargo.toml | 1 + .../xai-recsys-engine/src/request_metrics.rs | 98 +++++++++++-- .../xai-recsys-proto/proto/recsys.proto | 22 +++ phoenix/crates/storage/xai-o2/Cargo.toml | 16 +++ .../crates/storage/xai-o2/src/base_client.rs | 12 ++ phoenix/crates/storage/xai-o2/src/lib.rs | 7 + .../storage/xai-o2/src/o2_client_builder.rs | 115 +++++++++++++++ .../common/xai-proto/proto/recsys.proto | 22 +++ phoenix/reference/README.md | 2 +- phoenix/xrex/configs/data_feeds.py | 2 + phoenix/xrex/configs/xrecsys.py | 8 +- phoenix/xrex/configs/xrecsys_two_tower.py | 3 + phoenix/xrex/data/grpc_recsys.py | 2 + phoenix/xrex/data/parquet_recsys.py | 17 +++ phoenix/xrex/data/recsys/constants.py | 136 ++++++++++++++++++ phoenix/xrex/data/recsys/feature_config.py | 2 + phoenix/xrex/data/recsys/recsys_batch.py | 33 +++++ phoenix/xrex/data/recsys/sequence_packing.py | 14 +- phoenix/xrex/data/retrieval_dataset.py | 4 +- phoenix/xrex/data/streaming/kafkaloader.py | 14 +- phoenix/xrex/inference/launch_inference.py | 9 ++ phoenix/xrex/inference/model_runner.py | 44 +++++- phoenix/xrex/models/recsys_model.py | 100 ++++++++++++- phoenix/xrex/models/recsys_two_tower_model.py | 2 - phoenix/xrex/train/trainer_recsys.py | 79 +++++----- phoenix/xrex/utils/checkpoint_cloud.py | 4 +- phoenix/xrex/utils/metadata.py | 40 +++++- visibility-filtering/dark_traffic_setup.rs | 20 ++- 35 files changed, 856 insertions(+), 152 deletions(-) create mode 100644 phoenix/crates/storage/xai-o2/Cargo.toml create mode 100644 phoenix/crates/storage/xai-o2/src/base_client.rs create mode 100644 phoenix/crates/storage/xai-o2/src/lib.rs create mode 100644 phoenix/crates/storage/xai-o2/src/o2_client_builder.rs diff --git a/grox/config/config.py b/grox/config/config.py index 81e2065a..6f9575ea 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -88,7 +88,6 @@ class ModelName: EAPI_GROK_420_REASONING_INTERNAL = "eapi-grok-420-reasoning-internal" 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" diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index b353e40e..79743b06 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -74,14 +74,6 @@ def _fav_bucket(fav_count: int) -> str: half_open_max_calls=5, excluded_exceptions=(asyncio.CancelledError,), ) -_EAPI_4_5_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_6_INTERNAL_BREAKER_CONFIG = CircuitBreakerConfig( failure_rate_threshold=0.5, window_size=600.0, @@ -101,9 +93,6 @@ 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 ) @@ -111,8 +100,6 @@ def _fav_bucket(fav_count: int) -> str: ModelName.EAPI_GROK_4_6_INTERNAL, _EAPI_4_6_INTERNAL_BREAKER_CONFIG ) -_GROK_4_6_INTERNAL_DIAL = 0.3 - class SafetyPtosCategoryClassifier: result_pattern = re.compile(r"(.*)(.*)", re.DOTALL) @@ -200,8 +187,6 @@ class SafetyPtosChildSafetyPolicyClassifier: 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())) eapi_cfg_4_6 = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_6_INTERNAL) self.eapi_4_6_internal = EapiSampler( EapiModelConfig(**eapi_cfg_4_6.model_dump()) @@ -278,16 +263,10 @@ async def _cross_model_validate_with_4_5( ) -> SafetyPolicy: metric = "safety_ptos.child_safety_cross_model_validate_with_grok_4_5" try: - if random.random() < _GROK_4_6_INTERNAL_DIAL: - async with _eapi_4_6_internal_breaker.guard(): - raw = await self.eapi_4_6_internal.sample( - convo.interleaveToEapi(), conversation_id=convo.conversation_id - ) - else: - async with _eapi_4_5_internal_breaker.guard(): - raw = await self.eapi_4_5_internal.sample( - convo.interleaveToEapi(), conversation_id=convo.conversation_id - ) + async with _eapi_4_6_internal_breaker.guard(): + raw = await self.eapi_4_6_internal.sample( + convo.interleaveToEapi(), conversation_id=convo.conversation_id + ) confirm = self._parse_policy(raw) if confirm is None: logger.error( @@ -347,13 +326,6 @@ 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 - ) - self.eapi_4_5_internal = EapiSampler( - EapiModelConfig(**eapi_config_4_5_internal.model_dump()) - ) - eapi_config_4_6_internal = grox_config.get_eapi_model( ModelName.EAPI_GROK_4_6_INTERNAL ) @@ -474,12 +446,8 @@ async def classify_policy_for_violation( and violation.category in self.DELUXE_4_3_CATEGORIES ): if fav_count >= 1024: - if random.random() < _GROK_4_6_INTERNAL_DIAL: - mode = "deluxe-4.6-internal" - result = await self._sample_4_6_internal(convo) - else: - 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) @@ -543,31 +511,6 @@ 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 - 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( - 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" - ) - except Exception: - Metrics.counter("safety_ptos.eapi_4_5_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()}" - ) - return await self.llm.sample( - convo.interleave(), conversation_id=convo.conversation_id - ) - async def _sample_4_6_internal(self, convo: Conversation) -> str: breaker, sampler = _eapi_4_6_internal_breaker, self.eapi_4_6_internal try: diff --git a/grox/flows/reply_spam/task_filter.py b/grox/flows/reply_spam/task_filter.py index ee3518aa..5673574e 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 = 60000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 80000 @override @classmethod @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 60000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 80000 @override @classmethod 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 d807269e..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] @@ -78,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..1d5447c7 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 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/serving/xai-recsys-engine/Cargo.toml b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml index 99343fc9..91ecc659 100644 --- a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml +++ b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml @@ -64,6 +64,7 @@ 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 } 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..d36717f3 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs @@ -432,9 +432,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 { @@ -560,6 +558,7 @@ pub struct RequestMetricsGuard { admit_backlog: i64, admission: Option>, finished: bool, + recorded_reject: bool, } impl RequestMetricsGuard { @@ -571,9 +570,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 +607,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 +622,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 +639,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 +660,7 @@ pub fn on_request_finished(labels: &RequestLabels, result: &tonic::Result u64 { + NUM_REQUESTS_REJECTED + .with_label_values(&[reason, client]) + .get() + } + + #[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-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index efdb39f8..d274196b 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -226,6 +226,8 @@ message NextActionDistribution { map indexToContinuousValues = 7; RewardOutputs rewardOutputs = 8; + + SlateContext slateContext = 9; } message RewardOutputs { @@ -386,6 +388,26 @@ enum ActionName { 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_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; 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/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index efdb39f8..d274196b 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -226,6 +226,8 @@ message NextActionDistribution { map indexToContinuousValues = 7; RewardOutputs rewardOutputs = 8; + + SlateContext slateContext = 9; } message RewardOutputs { @@ -386,6 +388,26 @@ enum ActionName { 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_PURCHASE_CONVERSION = 200; ADS_MID_FUNNEL_CONVERSION = 201; ADS_ADD_TO_CART_CONVERSION = 202; 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 36e08cff..917ef4be 100644 --- a/phoenix/xrex/configs/data_feeds.py +++ b/phoenix/xrex/configs/data_feeds.py @@ -123,6 +123,7 @@ def _ranking_aggregated_kafka(mparams, hash_table, use_post_sid, sid_num_levels, 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", ""), ) @@ -148,6 +149,7 @@ def _ranking_rust_kafka(mparams, hash_table, use_post_sid, sid_num_levels, confi 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", ""), ) diff --git a/phoenix/xrex/configs/xrecsys.py b/phoenix/xrex/configs/xrecsys.py index 018d7e8c..3bd2deb5 100644 --- a/phoenix/xrex/configs/xrecsys.py +++ b/phoenix/xrex/configs/xrecsys.py @@ -416,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, @@ -638,10 +640,14 @@ 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), + 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"], diff --git a/phoenix/xrex/configs/xrecsys_two_tower.py b/phoenix/xrex/configs/xrecsys_two_tower.py index ea4875b0..453bf829 100644 --- a/phoenix/xrex/configs/xrecsys_two_tower.py +++ b/phoenix/xrex/configs/xrecsys_two_tower.py @@ -426,6 +426,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, diff --git a/phoenix/xrex/data/grpc_recsys.py b/phoenix/xrex/data/grpc_recsys.py index debf9b90..12615669 100644 --- a/phoenix/xrex/data/grpc_recsys.py +++ b/phoenix/xrex/data/grpc_recsys.py @@ -253,6 +253,7 @@ 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), + conversion_keep_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), @@ -369,6 +370,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), + conversion_keep_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 28150d59..09a1dc49 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 (_ckm := post_seq.get("conversion_keep_mask")) is not None: + padded["conversion_keep_mask"] = np.pad( + _ckm, ((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) @@ -1275,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 @@ -1347,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), + conversion_keep_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, @@ -1381,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 @@ -1476,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), + conversion_keep_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 17c7359c..ff350157 100644 --- a/phoenix/xrex/data/recsys/constants.py +++ b/phoenix/xrex/data/recsys/constants.py @@ -192,6 +192,7 @@ def to_pascal_case(s): "IsNotRelevantToSearch": [ "ClientTweetNotRelevantToSearch", ], + "IsSearchQueryReformulated": ["ClientTweetSearchQueryReformulated"], } ads_conversion_engagement_to_action_types = { @@ -224,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", ], @@ -251,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 + }, } @@ -261,6 +351,7 @@ def to_pascal_case(s): "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": {}, } @@ -285,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, @@ -296,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 0a287a2a..1c42d1f0 100644 --- a/phoenix/xrex/data/recsys/feature_config.py +++ b/phoenix/xrex/data/recsys/feature_config.py @@ -143,6 +143,8 @@ class UserInt64Feature(enum.IntEnum): "firstDpaProductKey", "authorFollowerCountSeq", "inReplyToPostIdSeq", + "is_delayed_feedback", + "conversionKeepMask", ] diff --git a/phoenix/xrex/data/recsys/recsys_batch.py b/phoenix/xrex/data/recsys/recsys_batch.py index 9e180030..f1ca8cf7 100644 --- a/phoenix/xrex/data/recsys/recsys_batch.py +++ b/phoenix/xrex/data/recsys/recsys_batch.py @@ -213,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 + conversion_keep_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 @@ -306,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] @@ -407,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) @@ -571,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_conversion_keep = 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( @@ -730,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_conversion_keep[*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, :] @@ -909,6 +922,7 @@ def _hash_dpa_keys(keys: npt.NDArray[np.int64], scale: int, bias: int) -> npt.ND product_surface=candidate_product_surface, client_app_id=candidate_client_app_id, post_ids=candidate_post_ids if include_candidate_post_ids else None, + conversion_keep_mask=candidate_conversion_keep, continuous_actions=candidate_continuous_actions, promoted_ids=candidate_promoted_ids, line_item_objective=candidate_line_item_objective, @@ -966,6 +980,15 @@ def _hash_dpa_keys(keys: npt.NDArray[np.int64], scale: int, bias: int) -> npt.ND 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) @@ -994,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"] + conversion_keep = post_seq.get("conversion_keep_mask") post_creation_ts_sec = post_seq["post_creation_ts_sec"] continuous_actions = post_seq["continuous_actions"] promoted_ids = post_seq["promoted_ids"] @@ -1043,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_conversion_keep = 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 ) @@ -1073,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 conversion_keep is not None: + new_conversion_keep[:, positive_slice] = conversion_keep 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 @@ -1239,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, + conversion_keep_mask=new_conversion_keep, continuous_actions=new_continuous_actions, promoted_ids=new_promoted_ids, line_item_objective=new_line_item_objective, @@ -1310,6 +1338,8 @@ def apply_global_negative_sampling( (batch_size, expanded_candidate_slots), dtype=client_app_id.dtype, ) + _gn_conversion_keep = post_seq.get("conversion_keep_mask") + new_gn_conversion_keep = 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, @@ -1343,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_conversion_keep is not None: + new_gn_conversion_keep[:, original_slice] = _gn_conversion_keep new_post_creation_ts_sec[:, original_slice] = post_creation_ts_sec new_continuous_actions[:, original_slice, :] = continuous_actions if categorical_features.shape[2] > 0: @@ -1466,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, + conversion_keep_mask=new_gn_conversion_keep, 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..aded39b0 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: + _ckm = batch["candidate_seq"].get("conversion_keep_mask") + assert _ckm is None or bool(np.asarray(_ckm).all()), ( + "conversion_keep_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 37335d5b..fb422f91 100644 --- a/phoenix/xrex/data/retrieval_dataset.py +++ b/phoenix/xrex/data/retrieval_dataset.py @@ -254,8 +254,8 @@ class RetrievalDataset(Enum): ) EVERGREEN = ( 5, - _idx("post_sid_v5_256x6_snapshots/evergreen_video_1825day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/evergreen_video_1825day.parquet"), + _idx("post_sid_v5_256x6_snapshots/video_4to14day.parquet"), + _idx("post_sid_v5_256x6_snapshots_backup/video_4to14day.parquet"), ) IMMERSIVENSFW = ( 6, diff --git a/phoenix/xrex/data/streaming/kafkaloader.py b/phoenix/xrex/data/streaming/kafkaloader.py index d2bec88d..3443fe56 100644 --- a/phoenix/xrex/data/streaming/kafkaloader.py +++ b/phoenix/xrex/data/streaming/kafkaloader.py @@ -759,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)) @@ -866,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}. " @@ -886,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: diff --git a/phoenix/xrex/inference/launch_inference.py b/phoenix/xrex/inference/launch_inference.py index 00f41e9e..d0cc1dc5 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: @@ -715,6 +716,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, diff --git a/phoenix/xrex/inference/model_runner.py b/phoenix/xrex/inference/model_runner.py index 1255922e..f5258b8b 100644 --- a/phoenix/xrex/inference/model_runner.py +++ b/phoenix/xrex/inference/model_runner.py @@ -166,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, @@ -634,6 +657,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) @@ -1537,6 +1561,21 @@ def _abort_hotswap_cycle(self, stage: str, reason: str) -> None: 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).") @@ -1636,6 +1675,7 @@ def _coordinator_loop(self) -> None: 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: @@ -1654,6 +1694,7 @@ def _coordinator_loop(self) -> None: self._swap_ready.set() self._swap_complete.wait() self._swap_complete.clear() + self._malloc_trim_after_swap() finally: self._hotswap_cycle_inflight.clear() @@ -1763,6 +1804,7 @@ def _follower_coordinator_loop(self) -> None: 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() @@ -4751,7 +4793,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), diff --git a/phoenix/xrex/models/recsys_model.py b/phoenix/xrex/models/recsys_model.py index 81b4a928..d784d97c 100644 --- a/phoenix/xrex/models/recsys_model.py +++ b/phoenix/xrex/models/recsys_model.py @@ -21,8 +21,11 @@ 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, ) @@ -525,6 +528,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 @@ -550,6 +555,12 @@ 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 @@ -1262,10 +1273,13 @@ def build_metric_masks( 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) @@ -1354,6 +1368,20 @@ def build_metric_masks( 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) + if condition_search_relevance_on_prompt: prompt_mask = jnp.any( raw_targets[:, :, jnp.array(SEARCH_RELEVANCE_ACTION_INDICES)] == 1, axis=-1 @@ -1367,6 +1395,15 @@ def build_metric_masks( 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 @@ -1555,6 +1592,7 @@ def _build_metric_masks( 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, @@ -1567,9 +1605,12 @@ def _build_metric_masks( 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( @@ -1585,6 +1626,7 @@ def compute_recsys_metrics( 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, @@ -1606,6 +1648,7 @@ def compute_recsys_metrics( line_item_objective, no_history_mask, dpa_product_key, + delayed_sample_mask, ) return self._compute_metrics_after_masks( @@ -2850,6 +2893,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 @@ -2994,6 +3054,11 @@ def loss( ) target_padding_mask = padding_mask[:, candidate_start_offset:] + conversion_keep_mask = batch["candidate_seq"].get("conversion_keep_mask") + if conversion_keep_mask is not None: + keep = cast_jax(conversion_keep_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[ :, @@ -3060,6 +3125,21 @@ 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 @@ -3069,6 +3149,15 @@ def loss( 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, @@ -3103,6 +3192,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, @@ -3116,6 +3206,7 @@ def loss( 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, @@ -3151,14 +3242,21 @@ def loss( 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_two_tower_model.py b/phoenix/xrex/models/recsys_two_tower_model.py index fb14083c..cbf51316 100644 --- a/phoenix/xrex/models/recsys_two_tower_model.py +++ b/phoenix/xrex/models/recsys_two_tower_model.py @@ -1582,8 +1582,6 @@ class RecsysTwoTowerModelConfig(Config): ads_only_candidates: bool = False - num_continuous_actions: int = 0 - multimodal_embedding_type: EmbeddingType | None = None user_features: UserFeaturesConfig = UserFeaturesConfig() diff --git a/phoenix/xrex/train/trainer_recsys.py b/phoenix/xrex/train/trainer_recsys.py index 78c942cc..28aa30de 100644 --- a/phoenix/xrex/train/trainer_recsys.py +++ b/phoenix/xrex/train/trainer_recsys.py @@ -603,6 +603,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 = { @@ -2332,6 +2334,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 @@ -3395,35 +3399,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(): @@ -3437,26 +3460,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/metadata.py b/phoenix/xrex/utils/metadata.py index 721edd87..fc856d10 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__) @@ -238,11 +239,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( @@ -290,13 +320,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 +345,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: diff --git a/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs index 5cb7486f..b4f97311 100644 --- a/visibility-filtering/dark_traffic_setup.rs +++ b/visibility-filtering/dark_traffic_setup.rs @@ -9,18 +9,24 @@ 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"; -const SHADOW_WORKLOAD: &str = "xai-vf-shadow"; +pub const STAGING_XDS_DEST: &str = "xai-vf-service.staging.visibility:grpc"; + +const FORWARDER_NAME: &str = "staging"; + +pub fn staging_tls_domain(dc: &str) -> String { + format!("visibility.visibility-filtering-service.staging.{dc}.s2s.twttr.net") +} pub type DarkLayer = Either; -struct StaticShadowDiscovery; +struct StaticStagingDiscovery; #[async_trait] -impl EndpointDiscovery for StaticShadowDiscovery { +impl EndpointDiscovery for StaticStagingDiscovery { async fn discover(&self) -> anyhow::Result> { Ok(vec![EndpointInfo { - name: SHADOW_WORKLOAD.to_string(), - xds_dest: format!("{SHADOW_WORKLOAD}.prod.visibility:grpc"), + name: FORWARDER_NAME.to_string(), + xds_dest: STAGING_XDS_DEST.to_string(), }]) } } @@ -50,7 +56,7 @@ pub fn resolve_layer() -> DarkLayer { } let dc = std::env::var("DATACENTER").unwrap_or_else(|_| "atla".to_string()); - let domain = format!("visibility.visibility-filtering-service.prod.{dc}.s2s.twttr.net"); + let domain = staging_tls_domain(&dc); info!(domain, "dark_traffic: enabled"); let factory = XdsChannelFactory::new( @@ -59,7 +65,7 @@ pub fn resolve_layer() -> DarkLayer { .with_domain_override(&domain), ); - let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(StaticShadowDiscovery)); + let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(StaticStagingDiscovery)); let config = ReloadableDarkTrafficConfigBuilder::new(CONFIG_PATH) .forwarders({ From d011592a1c8c4bfb23781ff15577a68dc08bdde1 Mon Sep 17 00:00:00 2001 From: CI agent Date: Mon, 24 Aug 2026 18:51:34 +0000 Subject: [PATCH 07/18] Open-source X Recommendation Algorithm --- grox/core/data_loaders/data_types.py | 24 + grox/core/data_loaders/post_mapper.py | 14 + grox/core/lm/post.py | 2 + grox/flows/ptos/classifier.py | 275 ++++-- grox/flows/ptos/task_safety_ptos_policy.py | 14 +- grox/flows/reply_spam/strato_loader.py | 17 + grox/flows/reply_spam/task_filter.py | 4 +- home-mixer/ads/util.rs | 5 +- .../ads_brand_safety_vf_hydrator.rs | 8 +- home-mixer/models/brand_safety.rs | 33 +- home-mixer/scored_posts_server.rs | 5 + .../ads_injection_logging_side_effect.rs | 8 +- .../xai_checkpointing/load.py | 214 ++++- visibility-filtering/config.rs | 5 + visibility-filtering/filter_tweets.rs | 36 +- visibility-filtering/lib.rs | 1 + visibility-filtering/reference_compare.rs | 812 ++++++++++++++++++ visibility-filtering/server_deps.rs | 51 +- visibility-filtering/twemcache/connection.rs | 7 +- 19 files changed, 1418 insertions(+), 117 deletions(-) create mode 100644 visibility-filtering/reference_compare.rs diff --git a/grox/core/data_loaders/data_types.py b/grox/core/data_loaders/data_types.py index b99e3045..0acd40eb 100644 --- a/grox/core/data_loaders/data_types.py +++ b/grox/core/data_loaders/data_types.py @@ -599,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 @@ -679,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] = [] @@ -788,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), @@ -822,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/post_mapper.py b/grox/core/data_loaders/post_mapper.py index e9b05973..aef93d11 100644 --- a/grox/core/data_loaders/post_mapper.py +++ b/grox/core/data_loaders/post_mapper.py @@ -15,6 +15,7 @@ ArticleMetadata as StratoArticleMetadata, ListMetadata as StratoListMetadata, ChatGroupMetadata as StratoChatGroupMetadata, + SpaceMetadata as StratoSpaceMetadata, ) from grox.core.data_loaders.data_types import ( Post, @@ -35,6 +36,7 @@ ArticleMetadata, ListMetadata, ChatGroupMetadata, + SpaceMetadata, AffiliatedBusiness, ) @@ -142,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, @@ -163,6 +170,7 @@ def from_post_metadata_strato(cls, post_metadata: StratoPostMetadata) -> Post: article_metadata=article_metadata, list_metadata=list_metadata, chat_group_metadata=chat_group_metadata, + space_metadata=space_metadata, ) @classmethod @@ -295,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 diff --git a/grox/core/lm/post.py b/grox/core/lm/post.py index e75b7cff..4842aee6 100644 --- a/grox/core/lm/post.py +++ b/grox/core/lm/post.py @@ -92,6 +92,8 @@ 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" diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index 79743b06..660491db 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -182,15 +182,231 @@ 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())) + + 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_5" + convo = self._build_policy_convo( + post, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, + illegal_and_regulated_behaviors_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"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_5_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_5_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_5_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_4_6 = grox_config.get_eapi_model(ModelName.EAPI_GROK_4_6_INTERNAL) - self.eapi_4_6_internal = EapiSampler( - EapiModelConfig(**eapi_cfg_4_6.model_dump()) - ) def build_convo(self, post: Post) -> Conversation: post_creation_time = ( @@ -216,10 +432,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: @@ -242,7 +454,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: @@ -252,49 +464,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.child_safety_cross_model_validate_with_grok_4_5" - try: - async with _eapi_4_6_internal_breaker.guard(): - raw = await self.eapi_4_6_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: diff --git a/grox/flows/ptos/task_safety_ptos_policy.py b/grox/flows/ptos/task_safety_ptos_policy.py index 83e9995e..efce237a 100644 --- a/grox/flows/ptos/task_safety_ptos_policy.py +++ b/grox/flows/ptos/task_safety_ptos_policy.py @@ -15,6 +15,7 @@ from grox.flows.ptos.classifier import ( SafetyPtosChildSafetyPolicyClassifier, SafetyPtosPolicyClassifier, + SafetyPtosPolicyCrossValidator, ) from grox.config.config import ModelName from grox.flows.ptos.mode import SafetyPtosMode @@ -28,6 +29,7 @@ class TaskSafetyPtosPolicyDetection(TaskWithPost): child_safety_classifier = SafetyPtosChildSafetyPolicyClassifier( gemma_model_name=GEMMA ) + cross_validator = SafetyPtosPolicyCrossValidator() classifiers = { SafetyPtosMode.STANDARD: SafetyPtosPolicyClassifier( @@ -91,8 +93,16 @@ 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 == SafetyPolicyCategory.ViolentMedia: + policy = await active_classifier.classify_policy_for_violation( + post, violation + ) + violation.safetyPolicy = await cls.cross_validator.validate( + violation.category, post, policy ) else: violation.safetyPolicy = ( diff --git a/grox/flows/reply_spam/strato_loader.py b/grox/flows/reply_spam/strato_loader.py index f882d985..e193fe51 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,9 +17,13 @@ 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 @@ -22,6 +31,14 @@ 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 5673574e..272a50a4 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 = 80000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 100000 @override @classmethod @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 80000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 100000 @override @classmethod diff --git a/home-mixer/ads/util.rs b/home-mixer/ads/util.rs index cc40aa7c..cb0e9f69 100644 --- a/home-mixer/ads/util.rs +++ b/home-mixer/ads/util.rs @@ -22,7 +22,10 @@ pub(crate) struct AdSpacing { } pub(crate) fn has_avoid(post: &ScoredPost) -> bool { - post.brand_safety_verdict() == BrandSafetyVerdict::MediumRisk + matches!( + post.brand_safety_verdict(), + BrandSafetyVerdict::MediumRisk | BrandSafetyVerdict::HighRisk + ) } pub(crate) fn find_safe_gaps(scored_posts: &[ScoredPost]) -> Vec { 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 6dc75c55..9bf776fc 100644 --- a/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs +++ b/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs @@ -102,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; } @@ -320,7 +320,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(); @@ -346,7 +346,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 @@ -503,7 +503,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/models/brand_safety.rs b/home-mixer/models/brand_safety.rs index 1ad616e7..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; } @@ -63,18 +72,13 @@ pub fn compute_verdict( } pub(crate) const MEDIUM_RISK_LABELS_V2: &[SafetyLabelType] = &[ - SafetyLabelType::NSFW_HIGH_PRECISION, - SafetyLabelType::NSFW_HIGH_RECALL, 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_V2, SafetyLabelType::GROK_NSFA_EXPANDED_V2, SafetyLabelType::NSFW_TEXT, @@ -99,6 +103,9 @@ pub(crate) fn compute_verdict_v2( 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; } @@ -175,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 ); } @@ -232,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, @@ -240,7 +247,7 @@ mod tests { ]); assert_eq!( compute_verdict(&labels, PRE_CUTOFF_ID), - BrandSafetyVerdict::MediumRisk + BrandSafetyVerdict::HighRisk ); } diff --git a/home-mixer/scored_posts_server.rs b/home-mixer/scored_posts_server.rs index 0c417b65..ff84fb1e 100644 --- a/home-mixer/scored_posts_server.rs +++ b/home-mixer/scored_posts_server.rs @@ -222,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/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/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py index 62fba013..09f526d4 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py @@ -1,8 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import contextlib +import ctypes +import fcntl import functools +import gc import json import logging +import math import os import pathlib import re @@ -26,6 +31,55 @@ PyTree = common.PyTree +_NODE_SERIALIZE_ENV = "XAI_RESTORE_NODE_SERIALIZE" +_NODE_LOCK_FILE_ENV = "XAI_RESTORE_NODE_LOCK_FILE" + + +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 _read_into_shards( t: ts.TensorStore, @@ -71,6 +125,7 @@ def load_checkpoint( domains: dict[str, Any] | None = None, tag: str | None = None, timeout: float = 900.0, + concurrent_gb: float | None = None, ): if tag is None: tag = "orbax-ckpt" @@ -105,8 +160,8 @@ def load_checkpoint( unloaded_state = host_state.copy() - futures = {} - for _i, checkpoint_name in enumerate(tree_to_dict(metadata, keep_none=False).keys()): + plan: list[tuple[str, str, list[bool], int]] = [] + for checkpoint_name in tree_to_dict(metadata, keep_none=False).keys(): name = checkpoint_name if rename is not None: name = rename(checkpoint_name) @@ -123,46 +178,133 @@ def load_checkpoint( 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) + 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))) del unloaded_state[name] - for future in common._ready(list(futures), timeout=timeout): - try: - future.result() - except Exception as e: - logger.exception(e) + 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: 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) + 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), + ) + + futures: dict[Any, tuple[str, str, int]] = {} + + def drain(): + 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 + 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 + rank_logger.debug("Loaded %s (name in checkpoint: %s)", name, checkpoint_name) + else: + rank_logger.debug("Loaded %s", 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, + ) - 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) + try: + 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: + 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) + 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), + ) + drain() + del stores + _release_batch_memory() + finally: + if node_lock is not None: + node_lock.close() rank_logger.info("Loading checkpoint took %.2f sec", time.time() - start) diff --git a/visibility-filtering/config.rs b/visibility-filtering/config.rs index 8646f3ab..2ddec299 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -3,6 +3,7 @@ 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_DUAL_CALL_HARNESS_ENABLED: &str = "VF_DUAL_CALL_HARNESS_ENABLED"; 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 = @@ -10,6 +11,10 @@ pub const ENV_MEDIA_FALLBACK_CACHE_SERVE_STALE_ENABLED: &str = pub const ENV_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED: &str = "VF_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED"; +pub fn dual_call_harness_enabled() -> bool { + parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) +} + pub fn fallback_cache_serve_stale_enabled() -> bool { parse_env_flag( std::env::var(ENV_FALLBACK_CACHE_SERVE_STALE_ENABLED) diff --git a/visibility-filtering/filter_tweets.rs b/visibility-filtering/filter_tweets.rs index aaa34654..1641ff31 100644 --- a/visibility-filtering/filter_tweets.rs +++ b/visibility-filtering/filter_tweets.rs @@ -1,7 +1,9 @@ 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::sync::Arc; use std::time::Instant; use tonic::{Request, Response, Status}; use tracing::info; @@ -9,11 +11,18 @@ 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( @@ -49,6 +58,15 @@ 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 { @@ -58,6 +76,20 @@ impl FilterTweetsEndpoint { 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() diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index 9809b07b..df8bc331 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -6,6 +6,7 @@ pub(crate) mod filter_tweets; pub(crate) mod get_safety_labels; pub mod hydration; pub mod models; +pub(crate) mod reference_compare; pub mod rules; pub mod safety_label_source; pub mod server; diff --git a/visibility-filtering/reference_compare.rs b/visibility-filtering/reference_compare.rs new file mode 100644 index 00000000..31ae39d4 --- /dev/null +++ b/visibility-filtering/reference_compare.rs @@ -0,0 +1,812 @@ +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 + }); + 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![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 && pages.last().is_some_and(|page| !page.is_empty()) { + pages.push(Vec::new()); + used = 0; + } + used += cost; + pages.last_mut().expect("pages is never empty").push(slice); + } + } + 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 = match reference_outcome { + Ok(results) => results, + 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) { + 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, + }; + + 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(None))).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/server_deps.rs b/visibility-filtering/server_deps.rs index c47c86ae..e035304c 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -4,6 +4,7 @@ use crate::filter_tweets::FilterTweetsEndpoint; use crate::get_safety_labels::GetSafetyLabelsEndpoint; use crate::hydration::{FallbackCacheMode, 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; @@ -18,6 +19,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"; @@ -190,6 +192,8 @@ 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; @@ -222,11 +226,49 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { ); 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> { + 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()) + ); + 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 TES_STRATO_REQUEST_TIMEOUT_MS: u64 = 100; fn tes_client_config(deterministic_aperture: bool) -> TESClientConfig { @@ -366,6 +408,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); diff --git a/visibility-filtering/twemcache/connection.rs b/visibility-filtering/twemcache/connection.rs index 150d5ce9..975f4e99 100644 --- a/visibility-filtering/twemcache/connection.rs +++ b/visibility-filtering/twemcache/connection.rs @@ -412,8 +412,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 +484,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}")))?; From 0d3cdd806c405f04db7030f720b48687aa304061 Mon Sep 17 00:00:00 2001 From: CI agent Date: Tue, 25 Aug 2026 23:20:04 +0000 Subject: [PATCH 08/18] Open-source X Recommendation Algorithm --- README.md | 6 +- .../service-lib/rules/enforcement_post.yaml | 7 +- .../service-lib/rules/enforcement_user.yaml | 17 +- .../service-lib/src/decision.rs | 1 + .../service-lib/src/facts.rs | 100 + .../service-lib/src/generic_actions.rs | 663 ++++ .../service-lib/src/growthbook.rs | 126 + .../service-lib/src/lib.rs | 295 +- .../service-lib/src/metrics.rs | 11 +- .../service-lib/src/rules.rs | 234 ++ grox/flows/reply_spam/generators.py | 2 +- grox/flows/reply_spam/plan_reply_ranking.py | 7 +- grox/flows/reply_spam/plan_spam_comment.py | 7 +- grox/flows/reply_spam/task_filter.py | 4 +- grox/flows/reply_spam/task_rate_limit.py | 10 - grox/flows/upa/constants.py | 2 + grox/flows/upa/generators.py | 11 + .../engagement_counts_hydrator.rs | 17 +- .../following_blocked_by_hydrator.rs | 65 + home-mixer/candidate_hydrators/mod.rs | 1 + .../phoenix_candidate_pipeline.rs | 5 +- .../reverse_chron_posts_pipeline.rs | 29 +- .../filters/brazil_2026_election_filter.rs | 3384 ++++++++++++++++- home-mixer/main.rs | 2 +- home-mixer/models/candidate.rs | 15 +- home-mixer/params/param.rs | 39 +- home-mixer/scorers/author_cold_start.rs | 21 +- home-mixer/scorers/ranking_scorer.rs | 229 +- home-mixer/scorers/vm_ranker.rs | 139 +- .../sources/following_night_owl_source.rs | 5 + home-mixer/util/author_rules.rs | 1 - home-mixer/util/urt/ad_marshaller.rs | 26 +- home-mixer/util/urt/client_event.rs | 34 +- home-mixer/util/urt/mod.rs | 21 +- .../util/urt/reverse_chron_following/mod.rs | 17 +- home-mixer/util/urt/wtf_marshaller.rs | 17 +- .../phoenixRankAllCandidateProcessor.strato | 15 + phoenix-rankall/src/config/mod.rs | 63 +- phoenix-rankall/src/store/base.rs | 263 +- phoenix-rankall/tests/integration_test.rs | 92 + .../xai-recsys-engine/src/copy_port_client.rs | 157 +- .../src/xai_configlib/__init__.py | 28 +- phoenix/xrex/driver/hooks.py | 35 +- phoenix/xrex/train/checkpoint_write.py | 2 + phoenix/xrex/train/trainer.py | 7 + phoenix/xrex/utils/metadata.py | 22 +- visibility-filtering/dark_traffic_setup.rs | 182 +- 47 files changed, 5850 insertions(+), 586 deletions(-) create mode 100644 abuse-enforcement-service/service-lib/src/generic_actions.rs create mode 100644 home-mixer/candidate_hydrators/following_blocked_by_hydrator.rs diff --git a/README.md b/README.md index fb83f9ca..e4c5656e 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ This repository contains the core code that determines which posts a viewer sees ## Table of Contents -- [Latest Updates](#latest-updates) +- [Notable Updates](#notable-updates) - [August 14th, 2026](#august-14th-2026) - [August 13th, 2026](#august-13th-2026) - [Overview](#overview) @@ -25,14 +25,14 @@ This repository contains the core code that determines which posts a viewer sees -## Latest Updates +## 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. 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). +- **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 25, 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 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..b2ab755c 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-25T16:15:48Z for_entity: user @@ -42,6 +42,10 @@ 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 +205,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/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/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 3055c182..e3060684 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; @@ -132,6 +133,7 @@ struct EnforcementCtx { rules_cache: Arc, allowlist: Option, kafka_producer_decisions: Option>, + kafka_producer_decisions_json: Option>, } struct ScoreResultProcessor { @@ -170,6 +172,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 +372,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 +443,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 +554,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 +572,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 +598,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 +638,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,7 +659,7 @@ 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) } @@ -663,7 +756,7 @@ impl ScoreResultProcessor { ]) .inc(); publish_decision_outcome( - self.ctx.kafka_producer_decisions.as_ref(), + &self.ctx, &decision_outcome( score, topic, @@ -716,7 +809,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()); @@ -1570,6 +1663,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(), }; { @@ -1874,12 +1968,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..2dd11592 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, } } @@ -1069,6 +1097,212 @@ 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_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/grox/flows/reply_spam/generators.py b/grox/flows/reply_spam/generators.py index c41828e9..5e5c080e 100644 --- a/grox/flows/reply_spam/generators.py +++ b/grox/flows/reply_spam/generators.py @@ -28,7 +28,7 @@ def _get_loader(self): @register class ReplyRankingRecoveryTaskGenerator(StreamTaskGenerator): TASK_GENERATOR_TYPE = REPLY_RANKING_RECOVERY - PLANS_TO_INJECT = {PlanReplyRanking.KEY} + PLANS_TO_INJECT = {PlanReplyRanking.KEY, PlanSpamComment.KEY} def _get_loader(self): return KafkaPostLoader(TOPIC_REPLY_RANKING_RECOVERY) 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/task_filter.py b/grox/flows/reply_spam/task_filter.py index 272a50a4..4565abfd 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 = 100000 + FOLLOWER_COUNT_THRESHOLD_FOR_SPAM_DETECTION = 120000 @override @classmethod @@ -182,7 +182,7 @@ async def _eligible_with_post(cls, post: Post, ctx: TaskContext) -> bool: class TaskReplyRankingFilter(TaskFilterWithPost): - FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 100000 + FOLLOWER_COUNT_THRESHOLD_FOR_REPLY_RANKING = 120000 @override @classmethod 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/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/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs b/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs index 0ef172d3..b87695f6 100644 --- a/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs +++ b/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs @@ -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() } @@ -102,6 +106,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 +165,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; } } @@ -221,7 +227,15 @@ mod tests { #[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 +246,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/mod.rs b/home-mixer/candidate_hydrators/mod.rs index 9796a9b6..27a4f3db 100644 --- a/home-mixer/candidate_hydrators/mod.rs +++ b/home-mixer/candidate_hydrators/mod.rs @@ -6,6 +6,7 @@ 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; diff --git a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs index c1c0c9cc..e8c9c397 100644 --- a/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs +++ b/home-mixer/candidate_pipeline/phoenix_candidate_pipeline.rs @@ -395,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]; diff --git a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs index 46a6da66..1ea8adf6 100644 --- a/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs +++ b/home-mixer/candidate_pipeline/reverse_chron_posts_pipeline.rs @@ -1,6 +1,7 @@ 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_following_candidate_hydrator::VFFollowingCandidateHydrator; @@ -8,6 +9,7 @@ use crate::clients::night_owl_client::{MockNightOwlClient, NightOwlClient, ProdN 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; @@ -20,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; @@ -50,6 +55,7 @@ impl ReverseChronPostsPipeline { strato_vf_client, xai_vf_client, vf_safety_labels_client, + socialgraph_client, ) = tokio::join!( async { Arc::new( @@ -94,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( @@ -102,6 +120,7 @@ impl ReverseChronPostsPipeline { strato_vf_client, xai_vf_client, vf_safety_labels_client, + socialgraph_client, ) .await } @@ -113,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 } @@ -123,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 { @@ -144,6 +165,7 @@ impl ReverseChronPostsPipeline { ]; let post_selection_hydrators: Vec>> = vec![ + Box::new(FollowingBlockedByHydrator::new(socialgraph_client).await), Box::new(VFFollowingCandidateHydrator::new( strato_vf_client, xai_vf_client, @@ -154,8 +176,11 @@ impl ReverseChronPostsPipeline { 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 index 2f832ca7..72f685c9 100644 --- a/home-mixer/filters/brazil_2026_election_filter.rs +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -7,25 +7,109 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // 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. +// 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. -// OmarAzizSenador deleted his account at the time this code was written. +// @OmarAzizSenador deleted his account at the time this code was written. +// @_ANDREDOPRADO no live account was found. +// @_EDUARDOMANTOAN no live account was found. +// @ADALBERTO_1111 no live account was found. +// @TWITTERADRIANAACCORSI no live account was found. +// @ADRIANASOUSAPIAUI no live account was found. +// @AGoldbach no live account was found. +// @AHELIXO no live account was found. +// @ALCEU_ALCEUMOREIRA no live account was found. +// @ALEXROSETI no live account was found. +// @ALKORAP1 no live account was found. +// @BetoRichaOficial no live account was found. +// @BRUNOPORTODEALMEIDA no live account was found. +// @CHARLES067277 no live account was found. +// @CRISTINAGRAEM no live account was found. +// @DANIELBRSOARES no live account was found. +// @DANILOBALASOFICIAL no live account was found. +// @DANILOTORRES100 no live account was found. +// @DECIOLIMAPT no live account was found. +// @DELEGADOEGUCHI no live account was found. +// @DEMAOLIVEIRA70 no live account was found. +// @DEPCELSOSABINO no live account was found. +// @DEPLUANAREGIA no live account was found. +// @DUARTEJR70 no live account was found. +// @DUDUSIVINSKI no live account was found. +// @EDSONSANTOSRJ no live account was found. +// @EUANGELAGARCIALINKTREE no live account was found. +// @EXPEDITOFUCAP no live account was found. +// @FADAPSICANALISE no live account was found. +// @FEDERALFELICIO no live account was found. +// @GERSONBURMANNIV no live account was found. +// @GSMA1986 no live account was found. +// @JAIZAMETODIO no live account was found. +// @LEOMASCARENHASP no live account was found. +// @LUCIANALIPPI30 no live account was found. +// @LUCIANAOROZIMBO no live account was found. +// @MARCELOSILVACAMPINAS no live account was found. +// @MEUCANA669499 no live account was found. +// @MIRCOCORONETTI no live account was found. +// @NADIAGERHARD no live account was found. +// @NETOFEITOSA6891 no live account was found. +// @PATRICIACRIZANTO2 no live account was found. +// @PAULOMOURAOTO no live account was found. +// @PEDRONASSIF_RJ no live account was found. +// @PEDROPONCIOBE no live account was found. +// @POLICIALPAULOBASTOS no live account was found. +// @SENATORCIDGOMES no live account was found. +// @XIGORPORTO no live account was 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([ + // @dayse + 6003262, + // @ricar + 6025402, + // @prado + 9171802, + // @madeleinelacsko + 9179462, // @renildo 14160928, // @renatoroseno 14492205, // @pedro_lupion 15022409, + // @soninhafrancine + 15768105, + // @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 @@ -40,36 +124,130 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 22480147, // @Pimenta13Br 22864100, + // @fabriziomeller + 22908605, // @rigotto 23443097, + // @gustavotutuca + 24055827, + // @camasao50 + 24103292, + // @radiovaldo + 24588898, // @euserafimcorrea 24761475, + // @rafangeli + 25164669, + // @pauloteixeira13 + 25562342, // @ManuelaDavila 25858078, + // @tomioyano + 26214420, // @CintyaMuniz 26284058, + // @Biango + 26560559, + // @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, + // @ticokuzma + 33759666, + // @CristinaMel + 33911227, // @aldenorlima 33973761, + // @_sidney_ + 34083807, + // @SharleneAZ + 34286009, + // @vimarchese + 34374755, + // @Peter_Costa + 34430921, + // @RafaCupertino + 34618485, + // @carlosviana + 34630924, + // @BetoRicha + 34665220, + // @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, + // @ruialves10_ + 37949658, + // @mauricioscalco + 38212742, + // @MarcioNakashima + 38585953, // @inacioarruda 38659819, + // @LucasCalilGo + 38929561, + // @tourinhopedro + 39818820, // @murilogaldinopb 40040727, // @FlavioBolsonaro @@ -78,46 +256,172 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 40463380, // @alexandrekalil 40721032, + // @RomanelliPR + 41271043, // @jessicamichels 41355867, + // @nidepp + 41403744, + // @EGCARREIRA + 41410587, + // @pauloabiackel + 41590536, + // @profcassiano + 42106436, + // @RaoniMendes + 42284334, // @lcbusato 42300711, // @Alice_Portugal 42304398, // @DepArthurMaia 42454330, + // @marciofrancasp + 42455446, // @kruke1 42487937, + // @ArlenSantiago + 42630237, + // @Francischini_ + 42736936, + // @onyxlorenzoni + 43041690, + // @pedropaulo + 43189774, + // @axelgrael + 43326346, + // @edurodrigues_25 + 43856097, + // @profleomatos + 44070650, // @marcelorangel1 44107208, + // @Fufaazevedo + 44153555, + // @renatasouzario + 44195955, + // @duarte_nogueira + 44455876, + // @rafaelcampelo + 44460118, // @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, + // @chrispuppi + 47418811, // @zeca_dirceu 47461491, + // @augustocury + 47529845, + // @naderaliumar + 47655960, + // @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, + // @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 @@ -128,56 +432,226 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 53050115, // @jorginhomello 53073647, + // @neyleprevost + 53163365, + // @renanroto + 53544192, + // @advandrebarros + 53681220, + // @AngeloAlmeidaBA + 53719510, + // @paulosalimmaluf + 53776600, + // @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, + // @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, + // @GustinhoRibeiro + 57208702, + // @Rafaelpicciani + 57529926, + // @CanzianiAlex + 57641073, + // @depHugoLeal + 57771926, + // @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, + // @marcio_motta + 60469505, + // @anapaulagold + 60708760, + // @claudioapolina + 60731692, + // @Isquierdorio + 60805457, + // @CovattiFilho + 60994156, // @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, + // @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, // @VerGuilherme @@ -188,68 +662,218 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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, + // @juniorsinforma + 69874322, // @Gyselle_Soares 70131422, // @ernanipolo 70284501, + // @pedroribeiropdt + 70301797, + // @janetepieta + 70361625, + // @lucarminatti + 70405671, + // @luis_fabuloso + 70443471, // @deputadoismael 70453020, + // @jonasdonizette_ + 70956209, + // @rdlorenzoni + 71056246, + // @charlesribeiro_ + 71098452, + // @lindberghfarias + 71310152, + // @lincolndrumond + 71541588, + // @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, + // @Casagrande_ES + 74722174, // @luizcoutopt 74738674, + // @VICENTINHOPT + 74762633, // @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, + // @michelschlemper + 80123403, // @capitaotadeu 80214491, + // @LuisTibeOficial + 80307396, + // @MarcioMacedoPT + 80548664, + // @aldairrizzi + 80559557, + // @heldersalomao + 80575628, + // @aniellefranco + 80582734, // @carlosmatosce 80597262, + // @EvairdeMelo + 80626542, // @depdarcidematos 80712669, // @Adrianageronim @@ -258,90 +882,264 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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, + // @giselecasarin + 84943825, + // @lubloureiro + 85150664, + // @RobertoPSOL + 85293890, + // @randolfeap + 85327394, + // @emanuelcacho + 85461555, + // @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, // @f_trad 92509126, + // @felipeaugusto01 + 92528438, + // @LucasRedecker + 92739280, + // @dpmarciomarinho + 92743190, + // @nancythame + 93065930, // @marcelodeputado 93116173, + // @bupessoa + 93424680, + // @danielsoranz + 93426187, + // @DuarteBechir + 93894534, // @tomazteixeira 93958167, + // @Obsevador + 94167900, // @coroneldavidms 94378206, // @ruycarneiropb 94428241, + // @netemoura + 94450303, // @samanthacavalca 94806323, + // @EdicarlosVieira + 94943878, // @requiaooficial 95253000, // @cabogilberto 95526088, + // @overissimo + 95621286, // @anapaulalimapt 95939603, // @PROFTULIO 96570084, + // @cassioandradepa + 96991211, + // @TrzeciakDaniel + 97587760, + // @evagoncalvesmg + 97677980, // @maxmacieldf 98786988, + // @matheusquintal + 98892109, + // @alexandrenau + 99001454, + // @ailtonlopespsol + 99389654, // @michelyfarina 99617723, + // @CacaLeao + 100521721, // @julioarcoverde 101016613, // @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, + // @CharlesDrumond + 109041140, + // @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, + // @ZeRicardoAM + 119079224, // @julio_cesar_pi 119115954, + // @capitaosamuelof + 119416561, // @adrianogaldino 119586707, + // @AndreQuintaoPT + 119818761, // @sibellebarros 120535860, // @aryvanazzi @@ -350,68 +1148,206 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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, + // @Danusalopes + 129028918, + // @apjunqueira + 129055364, + // @ninamarinabraga + 129070751, // @Mersinho_Lucena 129681031, + // @depleomonteiro + 129837652, // @alexandrebaldy 130620293, + // @abr + 131428902, + // @Miriampetrone + 132187525, // @depdelmasso 132190299, // @PedrodoOvo 132220469, // @Larissafgaspar 132480664, + // @betofantinel + 132735215, + // @40rodrigofarias + 133685785, // @glaustindafokus 133692726, + // @drthiagopeixoto + 134284069, // @sidneyleite_ 135349877, // @LucianoDucci 136004062, + // @matheuscadorin + 136381913, + // @paulomansur_ + 136710714, + // @BiradoPindare + 137548563, // @tadeuveneri 137919701, + // @Eduardo_Cury + 138504441, + // @Katiadiasjf + 139059131, // @MerlongSolano 140413929, + // @raphaelsebba + 140485170, + // @wiliantonezi + 141023529, + // @DiogoPBotelho + 141087783, + // @PSTUPE + 141090291, // @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, + // @RicardoCappelli + 153563550, + // @JarbasFilho_ + 155230905, // @dacassia1 155768056, + // @sgtalexandre + 156144539, // @profcanguru 156326262, + // @Sergio_Turra + 156355487, // @AzambujaReinald 157184730, + // @JoseMedeirosMT + 157226645, + // @RenataBuenoITA + 158067709, + // @Marivaldo4P + 159643822, + // @sandroalexpr + 160554168, // @MarcosRogerio 160895960, // @MarcoVamosaLuta 161318399, // @marinorpsol 161404367, + // @LianaCirne + 161416659, + // @maneco_hassen + 162794224, // @Vanderlan_VC 162807255, // @TJMFernandes 163251815, + // @emidinhomadeira + 163546697, + // @MariliaPFerrari + 163606193, + // @DepZeMilton + 163935204, + // @wandnogueira + 164012285, + // @dimasgadelha13 + 164058788, + // @DepNeilando + 164416842, // @EduardoBraga_AM 164439493, + // @profjosemarpsol + 164478764, + // @TiagoCado + 164816642, // @crisrbritto 165329128, + // @_Heloisa_Helena + 165499618, // @Jeronimoba13 165754961, // @DeputadoRoberio @@ -426,118 +1362,396 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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, + // @pepecollaco + 169158984, + // @neidermoreira + 169515284, + // @matheusmanholer + 170176086, // @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, + // @pablomarcal + 191223319, + // @lorran_rebeldia + 192010152, + // @esperidiaoamin_ + 192658045, + // @pretagilsa + 193045733, + // @analacerdamg + 193138273, + // @MarciaTaschetti + 198338329, + // @JacksonAndre7 + 198350574, + // @DepEduardoCunha + 198535390, + // @acmneto_ + 199025417, + // @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, + // @DanielBarbosaAL + 217322775, + // @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, + // @elzefacchinetti + 237562432, + // @wildermorais + 237687306, // @deltapericles 239036359, + // @_lucavalcante + 239059356, // @silascamara_ 243018634, // @UrsulaVidalPA 243429150, + // @gutoschiavetto + 243429651, + // @moarasaboia + 244491558, + // @depgurgel + 245320150, + // @RogerioCorreia_ + 245392082, // @DeAssisDiniz 247906787, + // @iginomarcos13 + 250070980, + // @prgilsondesouza + 250214822, // @ranypaulino 252553750, + // @RachelMaroja + 253219200, // @rdnarede 254172269, // @deprsantos 254201392, + // @SargentoFAHUR + 255300173, // @wellington_luiz 255637975, + // @depfrederico55 + 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, + // @pastorflamarion + 269618056, // @ayres_jr 269938072, + // @RafaelMottaRN + 271212223, // @alexgalvaodf 272477039, // @diegogarciapr 273616279, + // @JairMiotto + 273974254, // @doutorgutemberg 274515672, + // @BiaCerqueira_ + 274688443, + // @netocoelhoo + 276794985, + // @danealencar + 278126758, + // @paulodimelo + 278319124, + // @CostaMarinara + 278549268, + // @AllanPombopdt + 279092433, + // @InspetorRobison + 280306544, + // @mi_andrews + 284903425, + // @victorsalatiel + 285614672, + // @maicolmed + 286427700, + // @RenataMelorj + 286972716, + // @gabrielaorttiz + 287219048, // @marcelomaranata 287876093, + // @anaacarlinha + 288430204, + // @humberto130 + 288512944, // @delegado_waldir 288735634, + // @laualencar_ + 288761063, + // @ZeRobertoLula + 289318056, + // @waltercamargo40 + 289521136, + // @ArafetH + 290106695, // @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, + // @SamiraDaud + 310033093, + // @GildeteAlves + 310502840, + // @TercioTinoco + 311616488, + // @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, + // @cleciacarvalho1 + 334585581, + // @Fabinho_Gaspar + 334978230, + // @ZanaAmanda + 335065620, // @simboramudar22 337269106, + // @_soedi_ + 339854006, // @marisa_lobo 340331807, + // @danniellibrelon + 340748015, + // @MariliaFreireAM + 343512877, // @MarceloBelinati 345512946, // @edusantosdf @@ -546,124 +1760,360 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 349053302, // @paulinhoramosap 349163261, + // @DenianCouto + 350980280, + // @JanderBrum + 351008004, + // @paulapreta50 + 351272601, + // @RenatoLoffi + 351817280, + // @Haddad_Fernando + 354095556, + // @Biakicis + 357030742, + // @wanderleyporto + 358996581, // @andrefernm 360231929, + // @PROFAVILETE + 360649695, // @milkleileite 365210575, // @andresalineiro 367519089, // @Ana_claudiapb 370942708, + // @RONIESILVA15 + 373465468, + // @AdailtonAdvog + 379850268, + // @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, + // @luisfe_valdivia + 427885559, + // @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, + // @andreawerner_ + 475996406, + // @leopratesba + 480604517, + // @IanBlois + 480755252, + // @xerifedoconsum + 483269816, // @chicorafa_s 487108193, // @pedroluislongo 487622592, + // @LPescinelli + 492334281, + // @brunosouzasc + 494268633, + // @brauliolaranovo + 505339844, + // @glauberbastos_ + 527053144, // @MoisesSantosAc 538269241, // @JulianaBRedivo 545696318, + // @mariarosassp + 554803332, // @alcyvania 556852639, + // @ThammyReal + 578495086, // @OthelinoNeto 583377940, + // @HelioWirbiski + 589375704, // @AfonsoFlorence 599427558, + // @CiceroSimplicio + 604095777, + // @lucaspavanato + 608452346, + // @Jasson_Goulart + 608730390, + // @MarioEsteves2 + 610309440, + // @MateusWesp + 618506366, // @elmanooficial 626602522, // @fabiogov55 630758039, // @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, // @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, + // @AleSilvaOficial + 1089692132, // @depjanetedesa 1094959356, + // @NeumannJarbas + 1124876318, + // @McSmithOriginal + 1127018335, + // @cacateixeira45 + 1130957935, + // @daltonlueders + 1162799238, // @prof_juniorgeo 1226451780, // @marcoswesleymw 1308817182, // @alexceoficial 1314840228, + // @ThaysBieberbach + 1316758495, + // @zecarlospt + 1325494376, + // @bispadamares + 1326865753, + // @CARLOSVALADARE7 + 1356677952, // @D_GoretePereira 1362596354, + // @VAGNERVISOLI + 1420675674, + // @RobertoRocha_MA + 1436541721, // @mickasevalho 1461892208, + // @reginetebispo + 1485556436, + // @KimKataguiri + 1494658207, + // @leticiamattossc + 1498645730, // @deyvidbacelar 1526183576, + // @bublitz_a + 1532130170, + // @Luiza_RibeiroG + 1544201047, + // @natthpaccola + 1570635661, // @NegrahLima 1571332381, + // @oficialdmarques + 1604525460, + // @carlosfportinho + 1612753909, // @marciopachecopf 1613063005, + // @renancalheiros + 1650330319, // @GilvanMaximoOfc 1651852124, + // @carlosedrsantos + 1685646356, // @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, + // @ZeniteRosa + 2289590857, + // @BahMatteuss + 2293776768, + // @deltanmd + 2296138146, // @BrunoCarianha 2299610603, + // @MatSchilling + 2310607019, + // @delmartharocha + 2312487756, + // @ProfClaudioBran + 2352495202, // @DavidAlmeidaAM 2353403137, // @capitaoassis10 2359974485, + // @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, + // @profsoniameire + 2604536294, + // @depjorgesolla + 2605560932, + // @mvictoriabb + 2609799432, // @barbosinhams 2612421128, + // @_newtoncardoso + 2613918878, // @EderMauroPA 2632802395, // @paulo_litro @@ -672,114 +2122,436 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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, + // @_akalicia_ + 2879108776, // @paulolemosap 2881293743, // @manascimentogo 2892653500, // @LizianyM 2894163148, + // @RUBENSCANTUARIO + 2903026858, + // @andreicastroba + 2925491427, + // @KuhlmannJean + 2927136550, + // @DepJuscelino + 2970617333, + // @neioluciofp + 2977449981, + // @cabobonadiman + 2977624732, + // @guipaoficial + 2979670457, + // @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, + // @TeonilioBarba + 3119378914, + // @otavio_camp + 3125324049, + // @jarbassoaresjr + 3125532669, // @zuleidequeirozf 3130842411, + // @doutor_vicente + 3130887358, + // @DrLeonardomt + 3131609429, // @meire_cruvinel 3205786257, // @Marcio_Honaiser 3294107902, // @deputadopriante 3305760711, + // @marcos_reategui + 3306282352, + // @elikatakimoto + 3316550554, + // @ClaudiomirCast1 + 3335733075, + // @drjeanfreire + 3342622547, + // @PbnConcursos + 3357932231, // @moisesbrazpt 3373574517, // @kleybe_morais 3512854216, // @carmelonetobr 3662771592, + // @OPaiakan + 3674682197, + // @diegolopesadv + 3734258421, + // @delegadanadine + 3744465381, + // @heliomissao + 3853530796, + // @ninasouzarn + 3904595243, + // @FelipeMichelRJ + 4026638189, + // @geivissonvieira + 4028627062, + // @andersonrib_rj + 4031541093, // @paulo_mavignier 4043244676, + // @depsorayasantos + 4052741685, // @zeaugustonalin 4056653428, + // @GeraldoStocco + 4057530885, + // @Vinylobianco + 4077498658, + // @depzegeraldopt + 4204580127, // @_AlessandroSE 4250596815, + // @DepLucianoMDB + 4350613047, + // @josuelsantosbch + 4493572241, // @CarolDeToni 4566967516, + // @brunocunhablu + 4648157621, // @FofaBorges 4775264669, + // @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, + // @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, + // @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, + // @depeniotatto + 834065907743854592, + // @peumendonca23 + 839127759955841024, // @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, + // @brunopedralva + 889497158491271169, + // @jonesmanoel_PE + 889628930436730881, + // @MayconRobertoPR + 893177406705553410, + // @izadutrabr + 893178307591770112, + // @CatiaColombo2 + 894671613576269826, // @meuamigojoao 899602419302240257, // @todandara 904842960931565568, // @AlcyPinheiroCE 909049597091356673, + // @deborapsol + 915203945479458818, + // @leonidio_boucas + 917373635605786624, + // @cabojunioamaral + 917727546615193601, + // @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, + // @AnaFialho14 + 952537771717033984, + // @DpRicardoArruda + 953055428124045313, + // @Carlos_cabral81 + 956932819086913536, // @WilliamSiriRJ 958357145761845248, + // @delegadopiquet + 958833209961254912, + // @DepJorgeEverton + 959184931552464896, + // @Fbgg40 + 963928182368980993, + // @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 @@ -788,76 +2560,250 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 984221723544444928, // @eng_angelo44 984522534623301632, + // @AllanAguiar14 + 984581958943629313, + // @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, + // @prdinhosouza + 999747749364076545, + // @thiagoavilabr + 1000525959257378816, + // @Izalourenca + 1000724258740473856, + // @NatanSperafico + 1001476076466593792, + // @zuccors + 1002182052341534720, + // @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, + // @victorhugoforte + 1034054725107363840, // @RafagninLuciana 1034079788162535424, // @erikaamorimce 1034138653147242496, + // @AuricchioThiago + 1034858280231809024, + // @victoriogallimt + 1036259095458795526, // @AmauriRibeiroGO 1036681658760658945, + // @NeymarPesadao + 1037281805118910466, + // @FaleiroAirton + 1037384667836637185, + // @phbarroso45 + 1039361882845528064, + // @AlanaPassosRJ + 1039881241301016576, + // @DouglasGarcia + 1040704983358955526, // @RenanSantosMBL 1042601099566436352, + // @CavalarEmanuel + 1047079027696177152, + // @gutopfonseca + 1048016722433900550, + // @tome_abduch + 1049634819028717569, // @raphaelbarrabr 1049870508643233793, + // @ContaratoSenado + 1050121324436307970, + // @EuJorgeMiranda + 1051329992662142977, + // @Cristian0Engel + 1051904895106908161, + // @manumirella_ + 1052932498181840898, + // @depmgualberto + 1053123917785808901, + // @RafaelDemarchi5 + 1053334763858214912, + // @cleitinhotmj + 1057231251743170562, + // @delegadasheila + 1058010509256126464, + // @veronicalima_ve + 1059554967600685058, + // @pinheirinhomg + 1060134845043666945, // @pluviapt 1062505159824150530, // @capitaocarpe 1062678892216020992, + // @FlaviaHellen_13 + 1062871543007662080, + // @rafaelsimoesmg + 1064322572261699585, // @clarissatercio 1065379080084865024, + // @professorabebel + 1065745997391949824, + // @GIBERTOPINHEIRO + 1066799808172695552, + // @evandroaraujodf + 1069259896892329984, + // @DelHelioBressan + 1070740683386957825, + // @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, + // @BennyBriolly + 1084272914202091520, + // @majorfabianadep + 1084712593292443648, // @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 @@ -868,72 +2814,266 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1095754303367720960, // @depsargentolima 1095990374550695938, + // @DiogoTalento + 1096490764677320704, // @DFDanielFreitas 1097498693719199744, + // @dilvandafaroPT + 1097822283387809793, + // @joaocardosobr + 1097888011650502656, + // @thalesdacostaal + 1098031305818812416, + // @drmarcelmuscat + 1098089617209913344, // @CelAlfredo 1098328459825287187, + // @helio_ver + 1098911814824411137, // @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, + // @dredsondapaiol + 1110537499242389504, + // @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, + // @eunatashapoa + 1124140833547141121, + // @GilbertoCattan1 + 1125062220302385153, + // @Jeffers67876956 + 1125132737491480576, + // @RobertoJustusPR + 1125404580374822912, + // @FernandaLoubac2 + 1125416190199967744, + // @doutorfran + 1126519869053251584, + // @eduardo_borgo + 1126708527274315776, + // @diretoreliasjr + 1130259444326109184, + // @fabiosilvadep + 1130509226281975814, + // @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, + // @delegadopalumbo + 1179437585275465729, + // @rickazzevedo + 1184969015023800320, + // @PauloSussumu + 1188038521501749248, // @FranciscodoPT13 1188191474636206082, + // @samueljesusjb + 1188281892229074951, + // @atenabr051 + 1188827192999972865, + // @BrunoZambelli3 + 1191788412833075201, + // @ProtetorAle + 1193938478456868865, + // @thiagomedinamd + 1196894264334127107, + // @leninhamoc13 + 1197135895469711362, + // @JulianaBenicio_ + 1197281720694956032, + // @PettersNeto + 1198325133368340480, + // @DeputadoCarlosH + 1199293398194233344, + // @mariliacamposmg + 1199714008217063424, // @SandraLimadeVa1 1199740816018853890, + // @tukuma_pataxo + 1200840202194952194, + // @Ronygabriel_ofc + 1208038091962888194, // @LeoSuricate 1208544960032727040, + // @franciscodiasup + 1209429703909691392, // @ManuVieiraSC 1210296676520513536, // @CruzOrleans 1211689081861623808, + // @michelbeckerbr + 1211812977579421699, + // @mello_bandeira + 1213217041739419648, + // @mfriasoficial + 1213876635331465216, // @OficialNenemAl 1216746154626625536, + // @BrenoFonsecaMG + 1216897633769443329, // @JairSoutoAM 1218904655591243777, // @JohnRobertPA 1219003418347483136, + // @cortezpsol + 1219445803854548994, + // @mariamarighella + 1220387710462021633, + // @BernardinoNOVO + 1223018829393145856, + // @GuiBianco65 + 1223960824403939328, + // @VicMello16 + 1224035025882140673, + // @Vilsondafetaemg + 1224332353927073792, // @DrVictorAmoras 1224505061558161408, + // @luanlennonbr + 1224592244465905666, // @KlesleyGarcia 1224615230195609601, + // @joaobmaresguia + 1225055364124762112, // @profterezinhaPT 1225158156545904642, // @brenogaribalde @@ -944,90 +3084,240 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1226865113635926021, // @dramayraoficial 1226874382938787841, + // @marinahelenabr + 1227202002066837504, + // @VaniceMatos + 1227388838592614400, + // @JulioCLamim + 1227572253522526208, + // @Alfredogaspar_ + 1228008951813492739, // @BabaTupinamba 1232082766071767045, + // @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, + // @juliana_macieel + 1240694913420926976, + // @fabriciochaves_ + 1241145962082467851, + // @barcllay + 1241746085820932098, // @fabiolopes_38 1241821464111853570, + // @cop_santana + 1241842077979283466, // @elianexunakalo 1241878280841674752, + // @EiflerPam + 1241960263886286854, + // @depvermelho + 1242093483151839238, + // @claudiaguerramg + 1242834650155941888, // @samaramartinsup 1244783574676627460, + // @docporto + 1244954874049114113, + // @vitimporto + 1246125973755609088, + // @wmf_oficial + 1246641106391052288, + // @GallinatiRaquel + 1246945909881044994, // @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, + // @hana_ghassan + 1266449427985829888, + // @professorarita_ + 1266477339413749761, // @mariadoscamelos 1267127677468753920, + // @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, + // @ladisouzams + 1286438113896796170, // @ericodonovo 1287396715331452932, + // @sorriso_elisa + 1287490158510723072, + // @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 @@ -1036,88 +3326,352 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1310977778574032900, // @PabloSilvaLira 1311620519935045632, + // @gedalvaumbauba + 1312771507601498112, // @cassymonteiro 1312841615111843848, // @delegadkatarina 1314575606487670784, + // @andersonlimaadm + 1318194625400737792, + // @SerginhoCaxias + 1318204633823776770, + // @delegadalia + 1318672895573479424, + // @ElianaBayer_ + 1319073837485576197, // @anaportelams 1319825203313184768, // @DeoliAnderson 1322213572953493505, + // @PamelaGiedre + 1325859971905572869, + // @paulomelo_sa + 1330155785951866881, + // @cabomeireles + 1330393555290951681, + // @patisborges + 1330690808090124289, // @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, + // @SeccoHelio + 1355364623181094914, // @AnaPimentelmg 1356434859934298115, // @faustinorn01 1356762322489004032, + // @AndrePiresDF + 1358426024250343427, // @RafaelFonteles_ 1358739792163389440, + // @deputadodrPaulo + 1358880298407243778, + // @miroteixeira + 1359157620897153024, + // @FadelMoacyr + 1359162686228078593, + // @robertoclaudio + 1359194831415828488, + // @Nilvo17 + 1362188334924242949, + // @del_guilhermed + 1362442787988381696, + // @boulosnat + 1366796662614749187, + // @_mapeoficial + 1372519763818217477, + // @rodrigoestacho + 1374044251475091457, + // @PiauienseO + 1374436642375659530, // @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, + // @depmarcionunes + 1390475281316646917, + // @davibrandaobac1 + 1392451077656698881, + // @juarezcostamt + 1394317065750724608, + // @lucasbovesp + 1394364586346885120, + // @MassaelB + 1395784254269841408, + // @DelegadoMarcus + 1396055754793230341, + // @jusoaresft + 1400452986178981891, + // @oemersonmatos + 1400821681619353602, + // @eduacostario + 1400996132982038532, + // @antidio_lunelli + 1402602168838938625, + // @gu_camillo + 1405334928926162946, + // @Sandromarttinss + 1405740717595529225, + // @rubensuchoa_ + 1406043032437264386, + // @lelecopimentel_ + 1408810159376420865, + // @ananiasnauar + 1409194659872653317, // @sauloportolivei 1409385758838906882, // @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, + // @PeKelmon + 1437437148769226757, // @ScalcoDarlan 1437758990172184577, + // @nisia_trindade + 1438501569301991424, + // @Jadielmoraes20 + 1439280377936367621, + // @OperadoraManu + 1439763846772703232, + // @edilenxavier + 1441394605338021888, + // @mariiluse + 1445155253880594432, + // @renatmirandarj + 1446451214292553733, + // @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, // @josecam01577970 1459681826700673030, + // @schumarker7 + 1461714774996230145, // @RicardoArrruda 1462770465068437504, + // @BrunoBrazKart + 1463552472052502528, + // @KAUMAGN0 + 1463557962589429767, + // @Alexandregonrs + 1463574028468330497, + // @padovanidep + 1464329410115514369, + // @marinadomst + 1465395158103597065, + // @amarianalescano + 1468292579775348737, + // @FernandoManso20 + 1468757430632865793, + // @Maubmarcon + 1469007240279597067, + // @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, + // @EnfBrunoFarias + 1500910030593445888, + // @annasebbaj + 1501892430949453827, // @sgtgoncalves22 1504105286876991489, + // @EdianeMariaMTST + 1504228964319080449, + // @edsonferrazsc + 1505881785762304002, + // @AlvaroJeronymo + 1506029258967240715, // @SKaripuna 1506313161921777676, + // @rhdeverdade + 1508246534370082821, + // @Thais_ProfeChef + 1509245802815922177, + // @MaalouliMari + 1509589806325579784, // @orleansbrandao_ 1510790014111784967, + // @diegoandrademg + 1510954548348784645, + // @rodrigomarcial_ + 1511399052675633156, + // @WebaNatassia + 1512098373570113538, // @DelRodrigoSa 1512105284558282765, + // @MarciaHuculak + 1512756653199872006, // @reillerlopes 1513329845551390720, // @franze_carneiro @@ -1126,230 +3680,1050 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1515390927958908928, // @SocorroWaquiim 1516019212145287168, + // @PenhaBernardes + 1516033327941181447, // @DanielleDVale 1516121393028644869, + // @andrebuenoofc + 1517514873105797120, + // @ivanilsonrn + 1517576201338040321, + // @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, + // @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, + // @Gi_MonteiroRJ + 1553042102908502021, // @MissiasDias 1553167499901943808, + // @draalehaber + 1554462338715197446, + // @pedroaiharamg + 1555552838008414210, // @catanhopt 1555559206635315201, + // @PQueirozadv + 1555860014383988737, + // @LaironCarlos + 1556676552674394114, + // @KleberRosa50 + 1556777705387032577, // @VitormoreiraCG 1557086339740405761, + // @ScaranteRenato + 1557393886431027202, + // @Raphael89700864 + 1558100800664178688, + // @ArianeRSAssis + 1559215037746741251, + // @samuelvianamg + 1559598025420464134, // @Clecio_Luis 1560635958994812928, + // @BartonCutler + 1561212563836264449, + // @StellaGaio1224 + 1561545109715517440, + // @carlosmoraestv + 1563342678011842560, + // @vinivenades + 1564049349000183809, // @DeboraMenezes22 1564101654496133120, + // @drbenedettirs + 1567868352432951298, + // @NMousquer + 1569113308078170115, + // @pitmagrin + 1569391051160313858, + // @oAllaxSiqueira + 1569761874471972865, + // @regisethur + 1574862745325146112, // @weibetapeba 1575549757900132379, + // @14luizfranca + 1575672365094440960, + // @ProfAlexandreS + 1577470375272878083, // @Arnaldodeputado 1578127734362046489, + // @TamirisPeixoto1 + 1578450870127263757, + // @CarlosValdevin6 + 1579466590512418817, + // @eusouolimpio + 1580519722877128708, + // @quinhoprefeito + 1581818536137302021, + // @marlipaulinopr + 1583101243987206145, + // @renatodepaularj + 1583183342979190785, + // @WesleyCosta_GO + 1584241246817624064, + // @LucasLasmarMG + 1584270172952694784, + // @Joellobatoc + 1584299141550702593, + // @andreadantasc + 1584527735686397952, + // @rafandradembl + 1586192715137667072, + // @CombatPatriota + 1586416104632721416, + // @marcosfonsecapi + 1587131451593658369, // @antoniodoidoofc 1588344483766321154, + // @RONALDO90199231 + 1588685187629678592, // @gianninogueira2 1589336373101723649, + // @JUNIORCESARLEI9 + 1589412695664742401, + // @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, // @WickRyanAM 1612178912947183620, + // @SchiavoMaurilio + 1612183554787614722, // @deputadairacema 1612229099551858688, // @aveiltonsouza 1612484459034550273, + // @MalconMazzucato + 1612520249735135232, + // @depgilbertinho + 1612589448256016385, + // @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, // @AdrianaLeal2507 1637963109573828608, // @adrianalmeidapt 1638181881127612416, + // @leticiaaguiarsp + 1638274776413134849, + // @edercostarj + 1642933716279328768, + // @giuargolo + 1645831168586137600, + // @fialhooooo + 1646851212329754626, + // @Edy_Tocantins + 1647755717158264835, // @wistongomess 1648046380106104833, + // @CamilaGodoiSP + 1653466465855471617, + // @MatiasSamuka + 1655587131136307204, + // @DeputadoQuirino + 1656674946163277826, + // @DepComanDANte + 1657817923728166916, // @drbrunoresende_ 1661373663256408066, + // @CosttaIvony + 1662406778263420929, + // @wellersind + 1663242435420356609, // @DouglasRuas_RJ 1663563614954086401, // @bemoreiradf 1666457401816391682, + // @TathianaGuzella + 1668266165922144262, + // @carrarajh23 + 1669323757700169729, + // @LuizGracianoMBL + 1674888775506313216, + // @vanessarosajlle + 1675490861440802817, + // @GiFreitas1982 + 1684192295459905536, + // @victormenezesrj + 1686531661305966592, + // @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, + // @francisco_arten + 1710430261932908545, // @hebertcsgyn 1713936872370532352, + // @TitoBarichello + 1715811403708166144, + // @karisantospt + 1716415063471316993, + // @manubarrossp + 1717241739302371328, // @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, + // @samuel_al_silva + 1760365813456904192, + // @RicardoAlv32716 + 1761085576227213312, + // @FePresidente + 1762888285796343808, + // @dulce_lmendes + 1764371188489330689, + // @JeffreyChiquini + 1767171124629037056, // @Lenesilllva 1771346241449865216, + // @rafaelsatiebr + 1771625536772571136, + // @MayaraKeiko + 1773171382886576128, // @jorgemaciel05 1775273264774045697, + // @gisvaldopsol + 1776773968726298625, + // @Ap52467Jovelino + 1777297081134178304, + // @fabiocarneirojp + 1777538059891777537, + // @MarcoRo43166598 + 1777608369236271105, + // @igorrayanrn + 1777621309335183360, + // @juhliasantost + 1778077250908274688, + // @OficialVanucci + 1779252601269141504, + // @asargentolorena + 1779345802629914624, + // @carlosiranrs + 1779913435825795072, // @CamillaGonda 1781485235047174144, + // @brunnomattospt + 1782511140561367040, + // @MicheleMarxsc + 1782916900923543552, // @aucilene_a75929 1784209048671322112, // @eliabecamposs 1784275708606320641, + // @NivaldoNoga + 1785161935685562368, // @oescobarpro 1786921594390016001, + // @LeandroGol15952 + 1789440263187771392, + // @GuimaraesAlpha + 1790692252840243201, // @coronellrosses 1793428861348167680, + // @akarinaclaro + 1793823661369044992, + // @RicardoSeneseUP + 1795906221662208000, + // @nataliademesmao + 1797413322385485825, + // @Soldado_Sampaio + 1798332723989295104, + // @Lindenbergbra + 1799431977625366528, + // @maykondelfinomg + 1800001219932409856, + // @mqueiroz_rio11 + 1800605526683754496, + // @GilsonMachado22 + 1800983500541030400, + // @severoeulalio + 1801749578938564608, + // @leorondonn13 + 1803570674914426880, // @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, + // @drjoaomota832 + 1829618986171899905, + // @matheussimoespr + 1836487149769596928, + // @AGoldbach60024 + 1844153830201589760, + // @gabrielpiauhysp + 1844436268706418699, // @cironogueirapi 1844442355421614089, + // @ledapco + 1845277420854837249, + // @ProfMarcio58657 + 1845590017445400576, + // @GoulartVla23085 + 1846568278950346752, // @DragUrbana 1847431729415397376, + // @LuizaDoClezao + 1848153540012945408, // @abreu_de63729 1855074530613403648, + // @luladobemofici + 1855768833115541504, + // @debora__romani + 1855793756332519424, + // @LucasReis13_ + 1856394921189494784, // @sandsonmenezes 1856816774517268480, + // @Airtonjose26 + 1857173636689362944, + // @CristianeN74380 + 1858882502707810304, // @depcabomacielam 1859357204278542336, + // @denistaveiradn + 1865917852424798208, + // @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, // @ameliocayresdep 1887508211445702656, // @FelipeVasquesce 1889410539639517184, + // @celprincipebr + 1890456925147734016, + // @julianafideliis + 1896723491371581440, + // @gualbertoap + 1898078097452548097, + // @FRomulo13 + 1898340153896136704, // @Rpachecopinho 1898684990428151808, // @viscontioficial 1898913625697292288, // @GersonClaroMS 1899113980544897024, + // @galvaomicheles + 1899797969400180737, + // @Moanavaladares_ + 1899875749907357697, + // @caixeta_oficial + 1900205421123817472, + // @RafaMinatoSP + 1901680009430892544, + // @MacAntonioRJ + 1902731327461318656, + // @LenirOficial + 1904187270342574082, + // @profeVinicius + 1906160472354873344, + // @rodrigospada_ + 1908175828417949697, + // @MarinhoGui65411 + 1909642092424368128, // @glenioseixas 1909726014202171392, + // @yolandasilva_sc + 1911181511837011968, + // @betioldebochado + 1911652608915193856, + // @DrDiogoFranco + 1915836821826646016, + // @RenatoAngraRJ + 1916570073663475712, + // @aparecidobian + 1920469901136764928, + // @paulomeloparana + 1929325555234803712, + // @manoela__peres + 1932149081620484096, + // @leo_grandini13 + 1932924340917710848, // @prof_elson_sc 1937980838299504644, + // @marcioalvinosp + 1940410346411597824, + // @beto_vaz_ + 1944257271942287360, // @jotabrandaoam 1949664156606459905, + // @semeadorborges + 1952556730463707137, + // @DrSergioNeves + 1953095258784579586, + // @IgorMarquesBRA + 1957556152117653505, // @NoronhaPro7153 1958622002719182848, // @RubensAngiolett 1959013340661035008, + // @sebasticoelho + 1960040602877190144, + // @isabellaGedeon + 1960414219074998272, + // @DelAmirSalmen + 1961811802271952896, + // @blogmmedeiros + 1962570898310840320, + // @brunoboaretto + 1963656656342192129, + // @josearrudadf + 1964022075540312064, + // @juliemilkreal + 1965782186403012608, + // @delegadali47533 + 1966489240566444033, + // @edilsondamiaorr + 1970177461070528512, // @Drfilipecm 1972344363654037504, + // @romulobraz_13 + 1975026318975840257, // @DrManuelMarcos 1975317845962858501, + // @VereadorFox + 1975709711870906369, // @CarlosCostaPE10 1975873882654441472, + // @RangelJucy18837 + 1977431822557700096, + // @isapaixaosp + 1980463624649945088, + // @EliCorreaFilhoo + 1980794408141258752, + // @drjuliostobbe + 1981365553173299200, + // @mazzei_fel47870 + 1983280886067195904, // @marciaabrahaodf 1985760866109964288, + // @alepaivaoficial + 1986993348008419328, + // @GustavoHenryR + 1987721882838462464, // @DaClaudio33805 1988208205575712772, + // @RenatoBolsonar0 + 1988322911367950337, // @vandawitoto 1988323381792681984, // @SamaraMadureira 1990473255539621889, + // @YanProf + 1992656375382704128, + // @lemuel_SV + 1995470250738122752, + // @MarianaServente + 1996216267678875648, + // @jorgerrosario + 1997124812662362113, + // @oleandroissa + 1998066982999212032, + // @profakellsilva + 1998797409712009216, + // @nelsongrasselli + 2002006532184281088, + // @moisesbarboza + 2006965910309896192, + // @alvarenga35289 + 2007441310102523904, + // @AntoniadeJessp + 2010525266565857280, // @profraydf 2010751427623497728, + // @edinhosouzaaa + 2011825328793296896, + // @Efreu_Quintana + 2013273184640860160, // @brenomacedopi 2013660202873032707, + // @esthermoraessp + 2014421538418642944, // @robson_cacau 2015964533689323520, + // @carolontiveros0 + 2016185817392099336, + // @ingridcardososp + 2017575066381238272, // @sicchar1389 2018445464908042240, + // @vanesckaessusp + 2021065723573829633, + // @leninhavalente_ + 2021229327594000386, + // @MatheusCambuiBa + 2021559145728454656, + // @Jordambritosc + 2022354011035222016, + // @NetoFeitos68916 + 2026461466728292352, + // @catarinanevespb + 2026549782278529024, // @GreguyLoooban 2027415714953396224, // @AraceliLemosOF 2028870149546373120, + // @rfurtado22 + 2028873608903401472, + // @mendes__babi + 2029985728059559938, + // @helen_vitaRJ + 2033587304795963392, + // @EdneyBatalha + 2034242236678848513, + // @Brunodiasmissao + 2035542695536648192, + // @Fabio_x86 + 2036151026986696704, + // @MarcaoVivacqua + 2036163287587467265, + // @isamamedi027 + 2036512323020500992, // @missaogabi 2036817223558234112, + // @marinanamissao + 2036894494973440000, + // @isabeldesouzasp + 2036944910713081856, + // @stellabragasp + 2040479116068081664, + // @maurodeAL1930 + 2041451935685906432, + // @DepDrFlavio + 2041621533454450688, + // @vanessacfortes + 2043516611831701505, // @owilsonmartins 2043714770591735809, + // @viniciusdiaspi + 2045135821695574016, + // @CMolinariBR + 2046074615957540864, + // @VivianeSan11286 + 2046810859695947776, + // @emersonrrosa + 2047035719500103680, + // @Barbarabbotega + 2047324463612485632, + // @JoaoPaulo_2026 + 2047395841715920897, // @RitaDamore11 2047488478372368384, + // @VasconcellosCel + 2048607450593468416, + // @PablodoMST + 2048698490906169344, + // @CGasparin11011 + 2048961988999454720, + // @MarcioRezendeRJ + 2049539394701262848, + // @bia_pedagoga + 2049646600088117249, + // @oalanmansurrj + 2049927296442658816, // @jaimeverruckms 2052755712644730887, + // @Aminjhannouche + 2053565789421047808, // @ThefiAmancio21 2054040940885532672, + // @flaviadovalmg + 2054361311832616960, + // @marianasartor50 + 2054619154762674176, // @onenencoelho 2054729517864771584, + // @KeremhadassaMG + 2054936868026777600, + // @JulianaBrizola + 2054954070083813376, + // @mahmoudamer_rs + 2054974955603861504, + // @glaucelima12 + 2056410704031121408, // @edmartresoitao 2056559848750276608, + // @celmarcioasouza + 2056815030935416832, + // @RodolfoFiorucci + 2057456621865881600, + // @PL22Al + 2057808888297062400, + // @ProfNelsiWelter + 2058186878449278976, // @obrasilcomlula 2058595300873289728, + // @andreiamoura_df + 2059328849125490688, + // @eleusespaivaf + 2059680529746644993, // @vivianeluizams 2059726216119107587, // @ProfWiterNaves 2059790203091341316, + // @profcfabian + 2060000313843572736, + // @mayaelliz + 2060371890124857347, + // @cmidf_oficial + 2060549322207375368, + // @nicolasravipsol + 2061280727711289345, // @DrCrisVeloso 2061889166779019264, + // @NilsonVicentisc + 2062522303892611072, // @DaversonMatos 2062692988606717952, + // @jmonteirosc + 2062866727252250624, + // @CoronelMenezes_ + 2062867137497059328, + // @profjoaohs + 2063783489275600896, + // @aBiaAlcantara + 2064088166588440577, + // @DelmaPSOL + 2064731236023541760, + // @pr_itamar_paim + 2064819106545565696, + // @FelipeGambaroP + 2066731525400330240, + // @drcassiohprado + 2066997147971497984, + // @manoelsseverino + 2067256474024173568, + // @nenemalbuquerqu + 2067271467000045568, + // @SofiaFavero_ + 2067775247542095872, // @HVilelaporGoias 2069828777517973504, + // @verapinheirosc + 2069869784972369920, + // @Waltinhogo + 2070185289100775424, + // @Deborahzanchi_ + 2070650915871240192, + // @Delmariacorsato + 2070667303931256832, + // @transformarpsol + 2070866245293932545, + // @avictoriagallo + 2071779601815134208, // @angelaabukce 2071949683736354816, + // @Drrafujr + 2074177828606652416, + // @rogeriozabdalla + 2074246655759552513, + // @MariRibeiro42 + 2074876656561393664, + // @WiliansDouglaas + 2074899795404103681, // @drthalescoelho 2075213887213797376, + // @ArndaAcademia + 2075242788568899584, + // @RafaelPerlasca + 2076673278274383872, + // @AlinemunizRJ + 2076717583307255808, + // @zealexandrerj + 2076749095092256768, + // @CarineTAdv + 2077757274961911808, + // @LarianeTellMend + 2077832017736011777, // @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, + // @DanielSantanamc + 2081774169167974415, + // @luismartarj + 2081890844735426560, + // @drjoaomartins26 + 2082193801267888129, + // @Julio_Neto_SP + 2082200815763148800, + // @BibianoRN + 2082203221980778496, + // @DorvaniloNilo1 + 2082221380297236480, + // @SubTenSergio + 2082541411111493635, + // @anamariaa_of + 2082561982377406464, + // @Fabiani_vasco + 2082567723423342592, + // @Clarianabr + 2082581908005806081, + // @barbararesende0 + 2082924803145551872, + // @guihenriquesc + 2082932165466058752, + // @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, + // @panayotisdolula + 2085133776741437441, + // @diegojejees + 2085201272139886592, + // @lucianoleitoa_ + 2085736572914159616, + // @profandersonfig + 2085812265509388288, // @nandovianapsol 2086122224126136320, + // @boracomMarcel + 2086812805034815488, + // @roneymariachi + 2086873322587881473, + // @DarbideJesusrr + 2088266727519862784, + // @tatibarrapa + 2088334004445487104, + // @Mariguedes1406 + 2088832528383639552, ]) }); @@ -1537,7 +4911,7 @@ mod tests { #[test] fn hardcoded_list_is_non_empty() { assert!(!BRAZIL_2026_ELECTION_USER_IDS.is_empty()); - assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 665); + assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2328); } #[test] 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/candidate.rs b/home-mixer/models/candidate.rs index f6510138..3f2bfa2f 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -24,8 +24,6 @@ pub struct PostCandidate { pub slate_context: Option, #[serde(default)] pub served_slate_context: Option, - #[serde(default)] - pub mpn_parts: Option, #[serde( serialize_with = "serialize_served_type", deserialize_with = "deserialize_served_type" @@ -60,6 +58,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, @@ -94,6 +94,8 @@ pub struct SlateContext { pub sid_gap_l1: Option, pub sid_gap_l2: Option, pub sid_gap_l3: Option, + #[serde(default)] + pub recon_cos_milli: Option, } impl From for SlateContext { @@ -111,17 +113,11 @@ impl From for SlateContext { 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, } } } -#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct MpnParts { - pub pos: f64, - pub neg: f64, - pub scalar_multiplier: f64, -} - #[derive(Clone, Debug, Serialize, Deserialize)] pub struct SafetyLabelInfo { #[serde(with = "xai_safety_label_store::types::serde_label_type")] @@ -201,6 +197,7 @@ impl CandidateHelpers for PostCandidate { 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, }), reward_rerank_slot_prob: None, } diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index 845c06c7..eea74728 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-08-25T16:20:01Z 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, @@ -122,7 +122,7 @@ param!( PhoenixAggregationType, String, "rust_home_mixer_phoenix_aggregation_type", - "DENSE_WITH_SHORT_DWELL" + "DENSE_WITH_LONG_DWELL" ); param!( PhoenixRetrievalAggregationType, @@ -243,23 +243,16 @@ param!( "rust_home_mixer_log_slate_context", false ); -param!( - UseServedSlateContext, - bool, - "rust_home_mixer_use_served_slate_context", - false -); 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 ); @@ -336,7 +329,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); @@ -346,7 +339,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, @@ -360,7 +353,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, @@ -625,18 +618,6 @@ param!( "rust_home_mixer_vm_ranker_cluster_id", "Experiment3" ); -param!( - VMRankerValueModelId, - String, - "rust_home_mixer_vm_ranker_value_model_id", - "dpp" -); -param!( - VMRankerSendHeadWeights, - bool, - "rust_home_mixer_vm_ranker_send_head_weights", - false -); param!( VMRankerDppTheta, f64, @@ -715,7 +696,7 @@ param!( ColdStartTsTopK, u32, "rust_home_mixer_cold_start_ts_top_k", - 5 + 2 ); param!( ColdStartImpressionScale, @@ -928,7 +909,7 @@ param!( EnableAdsBrandSafetyVerdictV2, bool, "rust_home_mixer_ads_bs_v2_exp_enabled", - false + true ); param!( AdsTimeGapTSec, diff --git a/home-mixer/scorers/author_cold_start.rs b/home-mixer/scorers/author_cold_start.rs index 2340e28c..06c2f7f7 100644 --- a/home-mixer/scorers/author_cold_start.rs +++ b/home-mixer/scorers/author_cold_start.rs @@ -214,7 +214,7 @@ fn sample_reward( scale: f64, rng: &mut R, ) -> f64 { - let n = scale * candidate.view_count.unwrap_or(0) as 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); @@ -273,7 +273,7 @@ fn apply_cold_start( && cold_start_corpus_eligible(arm, c, corpus[*i]) && cold_start_freshness_eligible(arm, c, 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 < threshold) }) .map(|(i, _)| i) .collect(); @@ -411,7 +411,6 @@ rust_home_mixer: Arc::new(NullBucketImpressor::new()), None, false, - None, ) .unwrap(), ); @@ -428,32 +427,36 @@ 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 { - cold_start_candidate_with_favs(author_id, age, view_count, 0) + 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: u64, + 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 { 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), served_type: Some(pb::ServedType::ForYouPhoenixRetrievalMoe), ..Default::default() } diff --git a/home-mixer/scorers/ranking_scorer.rs b/home-mixer/scorers/ranking_scorer.rs index 962ab7ae..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( [ @@ -644,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() @@ -655,58 +570,18 @@ 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(); - let mut sid_counts: [FxHashMap; 3] = Default::default(); - let mut last_sid_rank: [FxHashMap; 3] = Default::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); - - let mut sid_k = [0u32; 3]; - let mut sid_gap = [None; 3]; - let sids = candidates[idx].semantic_ids.as_deref().unwrap_or(&[]); - let sid_known = !sids.is_empty(); - let mut prefix = 0u64; - for (level, &code) in sids.iter().take(3).enumerate() { - prefix = (prefix << 20) | (code as u32 as u64 & 0xFFFFF); - sid_k[level] = sid_counts[level].get(&prefix).copied().unwrap_or(0); - sid_gap[level] = last_sid_rank[level].get(&prefix).map(|last| rank - last); - sid_counts[level].insert(prefix, sid_k[level] + 1); - last_sid_rank[level].insert(prefix, rank); - } - - contexts[idx] = SlateContext { - k, - pool_rank: rank, - pool_rank_gap: rank_gap, - fatigue: 0.0, - pre_diversity_score: score, - sid_known, - sid_k_l1: sid_k[0], - sid_k_l2: sid_k[1], - sid_k_l3: sid_k[2], - sid_gap_l1: sid_gap[0], - sid_gap_l2: sid_gap[1], - sid_gap_l3: sid_gap[2], - }; + counts[idx] = k; author_counts.insert(author_id, k + 1); - last_author_rank.insert(author_id, rank); } - - contexts + counts } - fn served_slate_contexts( - query: &ScoredPostsQuery, - candidates: &[PostCandidate], - ) -> Option> { - if !query.params.get(UseServedSlateContext) { - return None; - } + fn served_slate_contexts(candidates: &[PostCandidate]) -> Option> { candidates.iter().map(|c| c.served_slate_context).collect() } @@ -714,25 +589,23 @@ impl RankingScorer { 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) @@ -800,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 @@ -815,53 +686,35 @@ impl Scorer for RankingScorer { None => false, }; - if mpn_scoring { - let persisted_contexts: Option> = - match Self::served_slate_contexts(query, candidates) { - Some(served) => Some(served), - None if query.has_cached_posts => Self::stored_slate_contexts(candidates), - None => 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) @@ -871,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() }) }) @@ -886,24 +734,8 @@ impl Scorer for RankingScorer { .author_cold_start .apply(query, candidates, &weighted_scores); - let persisted_contexts: Option> = - match Self::served_slate_contexts(query, candidates) { - Some(served) => Some(served), - None if query.has_cached_posts => Self::stored_slate_contexts(candidates), - None => 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() }; @@ -940,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; } } @@ -1013,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; @@ -1078,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; @@ -1099,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; @@ -1413,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 67207341..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,87 +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, - sid_known: s.sid_known, - sid_k1: s.sid_k_l1, - sid_k2: s.sid_k_l2, - sid_k3: s.sid_k_l3, - sid_gap1: s.sid_gap_l1, - sid_gap2: s.sid_gap_l2, - sid_gap3: s.sid_gap_l3, - }), - 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(); @@ -224,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/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/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 102584b1..3f646f37 100644 --- a/home-mixer/util/urt/mod.rs +++ b/home-mixer/util/urt/mod.rs @@ -88,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), @@ -259,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] @@ -267,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/reverse_chron_following/mod.rs b/home-mixer/util/urt/reverse_chron_following/mod.rs index 2a12825d..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; @@ -54,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 6a0aaa87..552e4fd2 100644 --- a/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato +++ b/phoenix-rankall-strato/columns/phoenix_rank_all/phoenixRankAllCandidateProcessor.strato @@ -260,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, @@ -452,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/src/config/mod.rs b/phoenix-rankall/src/config/mod.rs index d1f8db77..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,6 +178,10 @@ 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), @@ -183,7 +190,8 @@ impl PipelineKind { 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![], @@ -210,6 +218,7 @@ impl fmt::Display for PipelineKind { pub struct WindowConfig { pub name: String, pub retention: Duration, + pub min_age: Option, } impl WindowConfig { @@ -217,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), + } } } @@ -240,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" ); } @@ -254,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] @@ -334,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] 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/crates/serving/xai-recsys-engine/src/copy_port_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs index 4ae210e2..d977a84f 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 @@ -248,7 +248,11 @@ async fn run_downloads( .max(1); join_rate_limited(futures, limit, max_c).await } - _ => Ok(join_all(futures).await), + _ => 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}"))), } } @@ -268,7 +272,11 @@ async fn join_rate_limited( break; } let batch_size = batch.len(); - let batch_results = join_all(batch).await; + let batch_results = join_all(batch.into_iter().map(tokio::task::spawn)) + .await + .into_iter() + .collect::, _>>() + .map_err(|e| CopyPortError::Other(format!("copy_port download task join: {e}")))?; let failed = batch_results .iter() .filter(|r| r.0 == TRANSFER_FAILED_SENTINEL) @@ -1206,6 +1214,151 @@ 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_despite_spawn() { + 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], + "join_all on JoinHandles must keep submission order" + ); + assert_eq!(peak.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rate_limit_paces_between_spawned_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; expected ~1s pacing sleep" + ); + + 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/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/xrex/driver/hooks.py b/phoenix/xrex/driver/hooks.py index 7e05f432..fc8cee6a 100644 --- a/phoenix/xrex/driver/hooks.py +++ b/phoenix/xrex/driver/hooks.py @@ -584,8 +584,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 +597,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 +637,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 +675,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 +741,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 +752,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/train/checkpoint_write.py b/phoenix/xrex/train/checkpoint_write.py index 94e98152..ae5cf94f 100644 --- a/phoenix/xrex/train/checkpoint_write.py +++ b/phoenix/xrex/train/checkpoint_write.py @@ -84,6 +84,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 +97,7 @@ def _callback( elapsed_samples, self.current_ckpt_index, elapsed_tokens, + trainer_step, self.checkpoint_config.checkpoint_ttl, ) diff --git a/phoenix/xrex/train/trainer.py b/phoenix/xrex/train/trainer.py index ed7b4058..5df8d844 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 diff --git a/phoenix/xrex/utils/metadata.py b/phoenix/xrex/utils/metadata.py index fc856d10..e20a484d 100644 --- a/phoenix/xrex/utils/metadata.py +++ b/phoenix/xrex/utils/metadata.py @@ -33,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], ) @@ -57,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), ) @@ -67,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( @@ -75,6 +81,7 @@ def write_checkpoint_metadata( _CKPT_INDEX: checkpoint_index, _ELAPSED_TOKENS: elapsed_tokens, _CKPT_EXPIRY: checkpoint_expiry, + _STEP: step, }, f, ) @@ -93,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 @@ -105,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, @@ -112,6 +121,7 @@ def write_metadata_file( checkpoint_index, checkpoint_expiry, elapsed_tokens, + step, ) with (checkpoint_path / COMPLETED_FILENAME).open("x") as f: @@ -295,6 +305,7 @@ def record_checkpoint( checkpoint_index: int, checkpoint_ttl: int, elapsed_tokens: int, + step: int, ): ... @@ -428,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)) @@ -448,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) @@ -564,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 @@ -579,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/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs index b4f97311..4fbcc19d 100644 --- a/visibility-filtering/dark_traffic_setup.rs +++ b/visibility-filtering/dark_traffic_setup.rs @@ -1,17 +1,34 @@ use std::sync::Arc; -use tonic::async_trait; +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::{DynamicChannelManager, EndpointDiscovery, EndpointInfo}; +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 STAGING_XDS_DEST: &str = "xai-vf-service.staging.visibility:grpc"; -const FORWARDER_NAME: &str = "staging"; +pub const STAGING_NAMESPACE: &str = "visibility"; +pub const STAGING_APP_ENV: &str = "staging"; +pub const STAGING_PORT_ID: &str = "grpc"; +pub const STAGING_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") @@ -19,15 +36,115 @@ pub fn staging_tls_domain(dc: &str) -> String { pub type DarkLayer = Either; -struct StaticStagingDiscovery; +pub fn parse_staging_listener(listener_name: &str) -> Option { + let dest = listener_name.rsplit('/').next().unwrap_or(listener_name); + let dest = dest.split('?').next().unwrap_or(dest); + let suffix = format!(".{STAGING_APP_ENV}.{STAGING_NAMESPACE}:{STAGING_PORT_ID}"); + let workload = dest.strip_suffix(suffix.as_str())?; + if !workload.starts_with(STAGING_WORKLOAD_PREFIX) || workload.contains('.') { + return None; + } + Some(EndpointInfo { + name: workload.to_string(), + xds_dest: dest.to_string(), + }) +} + +struct XdsStagingDiscovery { + server_uri: String, +} -#[async_trait] -impl EndpointDiscovery for StaticStagingDiscovery { +#[async_trait::async_trait] +impl EndpointDiscovery for XdsStagingDiscovery { async fn discover(&self) -> anyhow::Result> { - Ok(vec![EndpointInfo { - name: FORWARDER_NAME.to_string(), - xds_dest: STAGING_XDS_DEST.to_string(), - }]) + tokio::time::timeout(LDS_RESPONSE_TIMEOUT, self.fetch()) + .await + .context("wildcard LDS exchange timed out")? + } +} + +impl XdsStagingDiscovery { + 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!("{STAGING_WORKLOAD_PREFIX}-dark-traffic"), + cluster: STAGING_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_staging_listener(name)) + .collect(); + + if endpoints.is_empty() { + tracing::warn!( + listeners = names.len(), + "dark_traffic: no staging listeners matched" + ); + } else { + info!( + names = %endpoints.iter().map(|e| e.name.as_str()).collect::>().join(", "), + "dark_traffic: discovery complete" + ); + } + + Ok(endpoints) + } +} + +struct TimeoutChannelFactory { + inner: XdsChannelFactory, +} + +#[async_trait::async_trait] +impl ChannelFactory for TimeoutChannelFactory { + async fn create_channel(&self, ep: &EndpointInfo) -> anyhow::Result { + tokio::time::timeout(CHANNEL_CREATE_TIMEOUT, self.inner.create_channel(ep)) + .await + .with_context(|| format!("channel dial timed out for {}", ep.xds_dest))? } } @@ -56,6 +173,9 @@ pub fn resolve_layer() -> DarkLayer { } let dc = std::env::var("DATACENTER").unwrap_or_else(|_| "atla".to_string()); + let discovery = XdsStagingDiscovery { + server_uri: format!("http://frontend.kube-discovery.prod.svc.{dc}.kube.int-x.ai:8082"), + }; let domain = staging_tls_domain(&dc); info!(domain, "dark_traffic: enabled"); @@ -65,7 +185,10 @@ pub fn resolve_layer() -> DarkLayer { .with_domain_override(&domain), ); - let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(StaticStagingDiscovery)); + let channels = DynamicChannelManager::new( + Arc::new(TimeoutChannelFactory { inner: factory }), + Arc::new(discovery), + ); let config = ReloadableDarkTrafficConfigBuilder::new(CONFIG_PATH) .forwarders({ @@ -103,6 +226,7 @@ mod tests { #[test] fn max_ordinal_threshold() { assert!(should_enable(Some(0), Some(3))); + assert!(should_enable(Some(1), Some(3))); assert!(should_enable(Some(2), Some(3))); assert!(!should_enable(Some(3), Some(3))); assert!(!should_enable(Some(4), Some(3))); @@ -112,4 +236,38 @@ mod tests { fn max_ordinal_zero_disables_all() { assert!(!should_enable(Some(0), Some(0))); } + + #[test] + fn parse_accepts_vf_staging_listeners() { + for (listener, workload) in [ + ("xai-vf-service.staging.visibility:grpc", "xai-vf-service"), + ( + "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_staging_listener(listener).expect(listener); + assert_eq!(ep.name, workload); + assert_eq!(ep.xds_dest, format!("{workload}.staging.visibility:grpc")); + } + } + + #[test] + fn parse_rejects_out_of_scope_listeners() { + for name in [ + "xai-vf-service.prod.visibility:grpc", + "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_staging_listener(name).is_none(), "{name}"); + } + } } From 45b48ba6baa40e212f6dcbaf8fe9fdc8d9da722e Mon Sep 17 00:00:00 2001 From: CI agent Date: Wed, 26 Aug 2026 20:04:12 +0000 Subject: [PATCH 09/18] Open-source X Recommendation Algorithm --- grox/core/lm/post.py | 24 +- grox/core/lm/thread.py | 29 +- .../reply_spam/classifier_reply_ranking.py | 1 + .../classifier_simple_reply_scorer.py | 4 +- grox/flows/reply_spam/strato_loader.py | 4 + grox/flows/reply_spam/task_write.py | 22 + home-mixer/models/candidate.rs | 8 + home-mixer/params/param.rs | 2 +- .../xai-recsys-engine/src/admission.rs | 19 + .../xai-recsys-engine/src/copy_port_client.rs | 31 +- .../xai-recsys-engine/src/request_metrics.rs | 30 ++ .../xai-recsys-proto/proto/recsys.proto | 12 + .../common/xai-proto/proto/recsys.proto | 12 + .../xai_checkpointing/load.py | 399 ++++++++++++++---- phoenix/xrex/configs/xrecsys_two_tower.py | 1 + phoenix/xrex/cutedsl/ranker_attention_fa4.py | 3 - .../cutedsl/ranker_attention_varlen_fa4.py | 3 - phoenix/xrex/data/grpc_recsys.py | 6 +- phoenix/xrex/data/parquet_recsys.py | 10 +- phoenix/xrex/data/recsys/recsys_batch.py | 28 +- phoenix/xrex/data/recsys/sequence_packing.py | 6 +- phoenix/xrex/inference/model_runner.py | 9 +- phoenix/xrex/models/recsys_model.py | 9 +- phoenix/xrex/models/recsys_two_tower_model.py | 13 +- phoenix/xrex/train/checkpoint_write.py | 1 + phoenix/xrex/train/misc.py | 4 + phoenix/xrex/train/trainer.py | 47 ++- phoenix/xrex/train/trainer_recsys.py | 48 ++- phoenix/xrex/utils/checkpointing.py | 221 +++++++++- phoenix/xrex/utils/metadata.py | 4 +- visibility-filtering/config.rs | 20 - visibility-filtering/dark_traffic_setup.rs | 57 ++- visibility-filtering/filter.rs | 1 - visibility-filtering/hydration/mod.rs | 5 +- .../hydration/tes_hydrator.rs | 296 +------------ visibility-filtering/server_deps.rs | 9 - 36 files changed, 873 insertions(+), 525 deletions(-) diff --git a/grox/core/lm/post.py b/grox/core/lm/post.py index 4842aee6..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: @@ -99,7 +105,13 @@ def render( 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/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/strato_loader.py b/grox/flows/reply_spam/strato_loader.py index e193fe51..dae31294 100644 --- a/grox/flows/reply_spam/strato_loader.py +++ b/grox/flows/reply_spam/strato_loader.py @@ -26,6 +26,10 @@ class ReplyRankingScoreStratoLoader: 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 diff --git a/grox/flows/reply_spam/task_write.py b/grox/flows/reply_spam/task_write.py index 1149af26..43d89d13 100644 --- a/grox/flows/reply_spam/task_write.py +++ b/grox/flows/reply_spam/task_write.py @@ -113,6 +113,28 @@ async def _publish_to_reply_ranking_manhattan( f"Missing user id [_publish_to_reply_ranking_manhattan] {reasoning=} {post.id=} {score=}" ) + existing = None + try: + existing = await ReplyRankingScoreStratoLoader.fetch_reply_ranking_score( + post.id + ) + except Exception: + logger.warning( + f"Failed to fetch existing reply ranking score for {post.id=}, proceeding with write" + ) + 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": "higher_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) diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index 3f2bfa2f..49d5fa88 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -96,6 +96,10 @@ pub struct SlateContext { 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, } impl From for SlateContext { @@ -114,6 +118,8 @@ impl From for SlateContext { 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, } } } @@ -198,6 +204,8 @@ impl CandidateHelpers for PostCandidate { 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, }), reward_rerank_slot_prob: None, } diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index eea74728..db2e30ee 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-25T16:20:01Z +// mirrored from config feature-switch defaults; last sync 2026-08-26T16:36:06Z use xai_feature_switches::param; param!( 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/copy_port_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs index d977a84f..c1e7b5cd 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 @@ -448,10 +448,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(); @@ -462,9 +458,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 prefix={prefix} bytes={bytes} in {:.2}s ({:.2} GB/s)", + secs, + gbs ); let prefix_slash = if prefix.ends_with('/') { @@ -485,10 +489,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, @@ -500,10 +500,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()) { 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 d36717f3..15f1aba7 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs @@ -452,6 +452,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, @@ -782,6 +788,30 @@ mod tests { .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"; diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index d274196b..5d69abef 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -110,6 +110,14 @@ message PredictNextActionsRequest { optional uint64 seedTweetId = 16; repeated uint32 requestedContinuousActionIndices = 17; + + repeated ConvAssetIds conv_asset_ids = 18; +} + +message ConvAssetIds { + repeated int64 impressed_time_ms = 1; + repeated int64 author_id = 2; + repeated int64 asset_id = 3; } message PredictNextActionsResponse { @@ -408,6 +416,9 @@ enum ActionName { 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; @@ -1240,6 +1251,7 @@ message SlateContext { optional uint32 sidGap1 = 10; optional uint32 sidGap2 = 11; optional uint32 sidGap3 = 12; + optional uint32 reconCosMilli = 13; } message ActionInfo { diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index d274196b..5d69abef 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -110,6 +110,14 @@ message PredictNextActionsRequest { optional uint64 seedTweetId = 16; repeated uint32 requestedContinuousActionIndices = 17; + + repeated ConvAssetIds conv_asset_ids = 18; +} + +message ConvAssetIds { + repeated int64 impressed_time_ms = 1; + repeated int64 author_id = 2; + repeated int64 asset_id = 3; } message PredictNextActionsResponse { @@ -408,6 +416,9 @@ enum ActionName { 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; @@ -1240,6 +1251,7 @@ message SlateContext { optional uint32 sidGap1 = 10; optional uint32 sidGap2 = 11; optional uint32 sidGap3 = 12; + optional uint32 reconCosMilli = 13; } message ActionInfo { diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py index 09f526d4..c10c7a38 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py @@ -26,7 +26,6 @@ import orbax.checkpoint as ocp -logger = logging.getLogger("checkpointing") rank_logger = logging.getLogger("rank") PyTree = common.PyTree @@ -81,6 +80,98 @@ def _release_batch_memory(): pass +def _build_read_plan( + metadata, + 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 tree_to_dict(metadata, keep_none=False).keys(): + 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, +) -> 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) + 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, array: jax.Array, @@ -117,6 +208,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], @@ -132,16 +260,7 @@ def load_checkpoint( 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) @@ -158,51 +277,13 @@ def load_checkpoint( } ) - unloaded_state = host_state.copy() - - plan: list[tuple[str, str, list[bool], int]] = [] - for checkpoint_name in tree_to_dict(metadata, keep_none=False).keys(): - 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))) - - del unloaded_state[name] + plan = _build_read_plan(metadata, 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: 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) + batches = _pack_read_batches(plan, concurrent_bytes) num_batches = len(batches) if concurrent_bytes: @@ -234,26 +315,6 @@ def load_checkpoint( futures: dict[Any, tuple[str, str, int]] = {} - def drain(): - 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 - - 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) - node_lock: _NodeBatchLock | None = None if _restore_node_serialize_enabled(): node_lock = _NodeBatchLock(_node_lock_path()) @@ -269,23 +330,15 @@ def drain(): with node_lock if node_lock is not None else contextlib.nullcontext(): stores = [] for checkpoint_name, name, mask, _nbytes in batch: - info = ocp.type_handlers.ParamInfo( - name=checkpoint_name, - path=path / checkpoint_name, - parent_dir=path, - is_ocdbt_checkpoint=True, - use_zarr3=use_zarr3, + t = _open_tensor( + checkpoint_name, + name, + path, + use_zarr3, + ts_context, + host_state[name], + has_domain=domains.get(name) is not None, ) - 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)) ): @@ -299,7 +352,7 @@ def drain(): inflight_bytes / (1 << 30), len(futures), ) - drain() + _drain_read_futures(futures, host_state, path, timeout, log_loaded=True) del stores _release_batch_memory() finally: @@ -309,6 +362,176 @@ def drain(): rank_logger.info("Loading checkpoint took %.2f sec", time.time() - start) +_AUTO_WINDOW_TARGET_BATCHES = 16 +_AUTO_WINDOW_MAX_BYTES = 32 * 10**9 + + +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, +) -> 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 + 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}, + } + ) + + plan = _build_read_plan(metadata, 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, + ) + + try: + 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, + ) + 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() + finally: + if node_lock is not None: + node_lock.close() + + rank_logger.info("Loading checkpoint (streamed) took %.2f sec", time.time() - start) + return replaced + + def broadcast_replicated( state: PyTree[jax.Array], axes: PyTree[jax.sharding.PartitionSpec], diff --git a/phoenix/xrex/configs/xrecsys_two_tower.py b/phoenix/xrex/configs/xrecsys_two_tower.py index 453bf829..0b5ce8d4 100644 --- a/phoenix/xrex/configs/xrecsys_two_tower.py +++ b/phoenix/xrex/configs/xrecsys_two_tower.py @@ -591,6 +591,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, diff --git a/phoenix/xrex/cutedsl/ranker_attention_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_fa4.py index 5240f25f..88bffaee 100644 --- a/phoenix/xrex/cutedsl/ranker_attention_fa4.py +++ b/phoenix/xrex/cutedsl/ranker_attention_fa4.py @@ -6,7 +6,6 @@ import jax import jax.numpy as jnp -from jax.ad_checkpoint import checkpoint_name _FA4_KERNEL_CACHE = {} @@ -362,8 +361,6 @@ def _attention(q, k, v): 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_bwd(res, g): diff --git a/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py index d850e1b1..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) @@ -670,8 +669,6 @@ 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) - out = checkpoint_name(out, "attn_outputs") - lse = checkpoint_name(lse, "attn_outputs") return out, (q, k, v, out, lse, bs_args) def _attention_bwd(res, g): diff --git a/phoenix/xrex/data/grpc_recsys.py b/phoenix/xrex/data/grpc_recsys.py index 12615669..1c26bc45 100644 --- a/phoenix/xrex/data/grpc_recsys.py +++ b/phoenix/xrex/data/grpc_recsys.py @@ -253,7 +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), - conversion_keep_mask=np.ones((batch_size, self.candidate_seq_len), dtype=np.bool_), + 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), @@ -370,7 +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), - conversion_keep_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), + 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 09a1dc49..f27721a4 100644 --- a/phoenix/xrex/data/parquet_recsys.py +++ b/phoenix/xrex/data/parquet_recsys.py @@ -768,9 +768,9 @@ def pad_array(arr: np.ndarray) -> np.ndarray: def pad_post_seq(post_seq: PostSeq) -> PostSeq: padded = _pad_post_seq_fields(post_seq) - if (_ckm := post_seq.get("conversion_keep_mask")) is not None: - padded["conversion_keep_mask"] = np.pad( - _ckm, ((0, batch_size - num_rows), (0, 0)), constant_values=True + 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 @@ -1361,7 +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), - conversion_keep_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), + 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, @@ -1492,7 +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), - conversion_keep_mask=np.ones((batch_size, candidate_seq_len), dtype=np.bool_), + 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/recsys_batch.py b/phoenix/xrex/data/recsys/recsys_batch.py index f1ca8cf7..efb2106a 100644 --- a/phoenix/xrex/data/recsys/recsys_batch.py +++ b/phoenix/xrex/data/recsys/recsys_batch.py @@ -213,7 +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 - conversion_keep_mask: NotRequired[npt.NDArray[np.bool_] | 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 @@ -581,7 +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_conversion_keep = np.ones(cand_shape_2d, dtype=np.bool_) + 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( @@ -742,7 +742,7 @@ def from_record_batch( candidate_product_surface[*cslice] = product_surface[*dslice] candidate_client_app_id[*cslice] = client_app_id[*dslice] if conversion_keep is not None: - candidate_conversion_keep[*cslice] = conversion_keep[*dslice] + 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, :] @@ -922,7 +922,7 @@ def _hash_dpa_keys(keys: npt.NDArray[np.int64], scale: int, bias: int) -> npt.ND product_surface=candidate_product_surface, client_app_id=candidate_client_app_id, post_ids=candidate_post_ids if include_candidate_post_ids else None, - conversion_keep_mask=candidate_conversion_keep, + trained_candidate_mask=candidate_trained_mask, continuous_actions=candidate_continuous_actions, promoted_ids=candidate_promoted_ids, line_item_objective=candidate_line_item_objective, @@ -1017,7 +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"] - conversion_keep = post_seq.get("conversion_keep_mask") + 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"] @@ -1067,7 +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_conversion_keep = np.ones((batch_size, total_candidate_slots), dtype=np.bool_) + 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 ) @@ -1098,8 +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 conversion_keep is not None: - new_conversion_keep[:, positive_slice] = conversion_keep + 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 @@ -1266,7 +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, - conversion_keep_mask=new_conversion_keep, + trained_candidate_mask=new_trained_mask, continuous_actions=new_continuous_actions, promoted_ids=new_promoted_ids, line_item_objective=new_line_item_objective, @@ -1338,8 +1338,8 @@ def apply_global_negative_sampling( (batch_size, expanded_candidate_slots), dtype=client_app_id.dtype, ) - _gn_conversion_keep = post_seq.get("conversion_keep_mask") - new_gn_conversion_keep = np.ones((batch_size, expanded_candidate_slots), dtype=np.bool_) + _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, @@ -1373,8 +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_conversion_keep is not None: - new_gn_conversion_keep[:, original_slice] = _gn_conversion_keep + 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: @@ -1498,7 +1498,7 @@ def apply_global_negative_sampling( product_surface=new_product_surface, client_app_id=new_client_app_id, post_ids=new_post_ids, - conversion_keep_mask=new_gn_conversion_keep, + 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 aded39b0..9c1eb6e0 100644 --- a/phoenix/xrex/data/recsys/sequence_packing.py +++ b/phoenix/xrex/data/recsys/sequence_packing.py @@ -166,9 +166,9 @@ def pack_batch( rng: np.random.Generator | None, block_size: int = 128, ) -> RecsysFeaturesBatch: - _ckm = batch["candidate_seq"].get("conversion_keep_mask") - assert _ckm is None or bool(np.asarray(_ckm).all()), ( - "conversion_keep_mask with masked candidates is not supported with sequence packing" + _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] diff --git a/phoenix/xrex/inference/model_runner.py b/phoenix/xrex/inference/model_runner.py index f5258b8b..e62ce54c 100644 --- a/phoenix/xrex/inference/model_runner.py +++ b/phoenix/xrex/inference/model_runner.py @@ -3253,7 +3253,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" ) @@ -4276,10 +4276,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( @@ -4290,7 +4293,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}", ) diff --git a/phoenix/xrex/models/recsys_model.py b/phoenix/xrex/models/recsys_model.py index d784d97c..3944837d 100644 --- a/phoenix/xrex/models/recsys_model.py +++ b/phoenix/xrex/models/recsys_model.py @@ -1381,6 +1381,9 @@ def build_metric_masks( 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( @@ -3054,9 +3057,9 @@ def loss( ) target_padding_mask = padding_mask[:, candidate_start_offset:] - conversion_keep_mask = batch["candidate_seq"].get("conversion_keep_mask") - if conversion_keep_mask is not None: - keep = cast_jax(conversion_keep_mask) + 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))) diff --git a/phoenix/xrex/models/recsys_two_tower_model.py b/phoenix/xrex/models/recsys_two_tower_model.py index cbf51316..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)) @@ -1259,7 +1266,9 @@ def _pool(outputs, mask, cu_seqlens): ) 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( @@ -1582,6 +1591,8 @@ class RecsysTwoTowerModelConfig(Config): ads_only_candidates: bool = False + use_history_segment_ids: bool = False + multimodal_embedding_type: EmbeddingType | None = None user_features: UserFeaturesConfig = UserFeaturesConfig() diff --git a/phoenix/xrex/train/checkpoint_write.py b/phoenix/xrex/train/checkpoint_write.py index ae5cf94f..6260ae50 100644 --- a/phoenix/xrex/train/checkpoint_write.py +++ b/phoenix/xrex/train/checkpoint_write.py @@ -126,6 +126,7 @@ 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, ) diff --git a/phoenix/xrex/train/misc.py b/phoenix/xrex/train/misc.py index 08aaac4e..87cbeb3b 100644 --- a/phoenix/xrex/train/misc.py +++ b/phoenix/xrex/train/misc.py @@ -53,6 +53,10 @@ class CheckpointConfig(Config): checkpoint_compressed: bool = True checkpoint_chunk_size_bytes: int = 1024 * 1024 * 4 + save_concurrent_gb: int | None = None + + restore_concurrent_gb: int | None = 32 + checkpoint_ttl: int = datetime.timedelta(weeks=2).total_seconds() replication_mode: Literal["full", "dp_only", "none"] = "dp_only" diff --git a/phoenix/xrex/train/trainer.py b/phoenix/xrex/train/trainer.py index 5df8d844..ea06718c 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -1160,6 +1160,12 @@ 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 maybe_load_checkpoint( self, ctx: TrainerContext, tag: str | None = None ) -> tuple[bool, int, int]: @@ -1167,13 +1173,30 @@ def maybe_load_checkpoint( rank_logger.info("Not loading checkpoint; starting from scratch") return False, 0, 0 + do_not_load_opt_state = ( + self.checkpoint_config.no_opt_state or self.reinit_on_load + ) and ctx.checkpoint.is_manual_load() + 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 pinned-host staging)" + ) + self.host_state = jax.device_put( + warm_purge(self.state), + warm_purge(restore_staging_sharding), + ) + else: + do_not_load_opt_state = False + self.host_state = jax.device_put(self.state, restore_staging_sharding) rename = None @@ -1182,10 +1205,6 @@ def maybe_load_checkpoint( 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 do_not_load_opt_state: host_state = self.purge_opt_state_on_load(host_state) @@ -1238,11 +1257,23 @@ def maybe_load_checkpoint( 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 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 = {} diff --git a/phoenix/xrex/train/trainer_recsys.py b/phoenix/xrex/train/trainer_recsys.py index 28aa30de..2c53a37f 100644 --- a/phoenix/xrex/train/trainer_recsys.py +++ b/phoenix/xrex/train/trainer_recsys.py @@ -1331,8 +1331,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) @@ -1343,7 +1353,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"], @@ -1357,15 +1367,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, @@ -1373,7 +1374,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, ) @@ -1402,7 +1403,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, {} @@ -1512,8 +1513,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) @@ -1522,7 +1527,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"], @@ -1530,6 +1535,8 @@ 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: @@ -1544,10 +1551,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( @@ -1555,7 +1561,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( @@ -1712,7 +1718,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] @@ -2439,6 +2444,11 @@ def purge_opt_state_on_load(self, host_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) diff --git a/phoenix/xrex/utils/checkpointing.py b/phoenix/xrex/utils/checkpointing.py index 5047ae4e..cd287d22 100644 --- a/phoenix/xrex/utils/checkpointing.py +++ b/phoenix/xrex/utils/checkpointing.py @@ -1,7 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import contextlib +import ctypes +import fcntl +import gc import logging import os +import time +from typing import Any import jax import jax.numpy as jnp @@ -16,16 +22,87 @@ 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 + + +def get_checkpointer(timeout_secs=900, save_concurrent_gb: int | None = None): + global _CHECKPOINTER, _CHECKPOINTER_SAVE_CONCURRENT_GB if _CHECKPOINTER is None: + 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 = ocp.AsyncCheckpointer( - ocp.PyTreeCheckpointHandler(use_zarr3=True), timeout_secs + 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 +115,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 +221,7 @@ def save_checkpoint( compressed=True, chunk_byte_size: int = 1024 * 1024 * 4, tracer: Tracer | None = None, + save_concurrent_gb: int | None = None, ): if not tracer: tracer = MockTracer() @@ -57,7 +230,7 @@ def save_checkpoint( tag = "orbax-ckpt" os.makedirs(path, exist_ok=True) - checkpointer = get_checkpointer(timeout_secs) + checkpointer = get_checkpointer(timeout_secs, save_concurrent_gb=save_concurrent_gb) with tracer.start_as_current_span("wait_for_previous_checkpoint"): try: @@ -86,10 +259,34 @@ 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: + 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) @@ -113,10 +310,22 @@ def _callback(): ), ) - 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") diff --git a/phoenix/xrex/utils/metadata.py b/phoenix/xrex/utils/metadata.py index e20a484d..c4920ef9 100644 --- a/phoenix/xrex/utils/metadata.py +++ b/phoenix/xrex/utils/metadata.py @@ -254,7 +254,7 @@ def guess_checkpoint_format(path): raise ValueError(f"Could not determine format of checkpoint at {path}") -def _has_committed_payload(checkpoint_path: Path) -> bool: +def has_committed_payload(checkpoint_path: Path) -> bool: if (checkpoint_path / "ckpt-0" / "tensor00000_000").exists(): return True final_names = set() @@ -272,7 +272,7 @@ def _has_committed_payload(checkpoint_path: Path) -> bool: def _is_loadable_checkpoint(checkpoint_path: Path) -> bool: if not (checkpoint_path / COMPLETED_FILENAME).exists(): return False - if _has_committed_payload(checkpoint_path): + 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" diff --git a/visibility-filtering/config.rs b/visibility-filtering/config.rs index 2ddec299..1eff08b3 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -6,10 +6,6 @@ pub const ENV_GRPC_MTLS_CLIENT_CA_PATH: &str = "GRPC_MTLS_CLIENT_CA_PATH"; pub const ENV_DUAL_CALL_HARNESS_ENABLED: &str = "VF_DUAL_CALL_HARNESS_ENABLED"; 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 fn dual_call_harness_enabled() -> bool { parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) @@ -31,22 +27,6 @@ pub fn fallback_cache_populate_enabled() -> bool { ) } -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 media_fallback_cache_populate_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_MEDIA_FALLBACK_CACHE_POPULATE_ENABLED) - .ok() - .as_deref(), - ) -} - fn parse_env_flag(value: Option<&str>) -> bool { value.is_some_and(|value| { matches!( diff --git a/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs index 4fbcc19d..8626a983 100644 --- a/visibility-filtering/dark_traffic_setup.rs +++ b/visibility-filtering/dark_traffic_setup.rs @@ -157,17 +157,19 @@ pub fn resolve_layer() -> DarkLayer { 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_enable(ordinal, max_ordinal) { + if !should_mirror(workload.as_deref(), ordinal, max_ordinal) { info!( + ?workload, ?ordinal, ?max_ordinal, - "dark_traffic: disabled (ordinal >= max)" + "dark_traffic: disabled (not a mirror host)" ); return Either::Right(tower::layer::util::Identity::new()); } @@ -200,41 +202,62 @@ pub fn resolve_layer() -> DarkLayer { Either::Left(DarkTrafficLayer::new(config)) } -fn should_enable(ordinal: Option, max_ordinal: Option) -> bool { - let max = max_ordinal.unwrap_or(1); - let ord = ordinal.unwrap_or(u32::MAX); - ord < max +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_enable(Some(0), None)); - assert!(!should_enable(Some(1), None)); - assert!(!should_enable(Some(99), None)); + 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_enable(None, None)); - assert!(!should_enable(None, Some(3))); + assert!(!should_mirror(PROD, None, None)); + assert!(!should_mirror(PROD, None, Some(3))); } #[test] fn max_ordinal_threshold() { - assert!(should_enable(Some(0), Some(3))); - assert!(should_enable(Some(1), Some(3))); - assert!(should_enable(Some(2), Some(3))); - assert!(!should_enable(Some(3), Some(3))); - assert!(!should_enable(Some(4), Some(3))); + 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_enable(Some(0), Some(0))); + 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] diff --git a/visibility-filtering/filter.rs b/visibility-filtering/filter.rs index c2d6bb82..18aeaa95 100644 --- a/visibility-filtering/filter.rs +++ b/visibility-filtering/filter.rs @@ -187,7 +187,6 @@ mod tests { socialgraph, labels, crate::hydration::FallbackCacheMode::Disabled, - crate::hydration::FallbackCacheMode::Disabled, ), Policies::new(), ) diff --git a/visibility-filtering/hydration/mod.rs b/visibility-filtering/hydration/mod.rs index ff66570c..33e088c0 100644 --- a/visibility-filtering/hydration/mod.rs +++ b/visibility-filtering/hydration/mod.rs @@ -136,13 +136,14 @@ impl HydrationPipeline { socialgraph_client: Arc, safety_label_source: Arc, fallback_cache_mode: FallbackCacheMode, - media_fallback_cache_mode: FallbackCacheMode, ) -> 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, diff --git a/visibility-filtering/hydration/tes_hydrator.rs b/visibility-filtering/hydration/tes_hydrator.rs index 35ff5a93..cf873fd5 100644 --- a/visibility-filtering/hydration/tes_hydrator.rs +++ b/visibility-filtering/hydration/tes_hydrator.rs @@ -1,5 +1,4 @@ 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, @@ -14,11 +13,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)] @@ -47,16 +44,6 @@ impl TweetHydration { } 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 +75,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(); @@ -171,14 +154,6 @@ 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), @@ -187,7 +162,7 @@ impl TesHydrator { 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), } } @@ -284,10 +259,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 +290,9 @@ mod tests { } fn hydrator() -> TesHydrator { - TesHydrator::new( - Arc::new(MockTESClient::default()), - FallbackCacheMode::Disabled, - ) + TesHydrator { + tes_client: Arc::new(MockTESClient::default()), + } } #[test] @@ -512,263 +483,4 @@ mod tests { 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_trusted_friends_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/server_deps.rs b/visibility-filtering/server_deps.rs index e035304c..2da74a1c 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -98,13 +98,6 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { } 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 tes_client: Arc< dyn xai_core_entities::tweet_entity_service_client::TESClient + Send + Sync, @@ -208,7 +201,6 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { sg_client, safety_label_source.clone(), fallback_cache_mode, - media_fallback_cache_mode, ); let policies = crate::rules::Policies::new(); let (home_rule_count, recommendations_rule_count) = policies.rule_counts(); @@ -219,7 +211,6 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { info!( hydrator_count = 5, ?fallback_cache_mode, - ?media_fallback_cache_mode, home_rule_count, recommendations_rule_count, "VFServer initialized with prod clients" From 24c60942c5c5fdad3a6addffb4c6e6d2f228f04f Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 28 Aug 2026 01:00:41 +0000 Subject: [PATCH 10/18] Open-source X Recommendation Algorithm --- .../service-lib/rules/enforcement_user.yaml | 4 +- .../service-lib/src/ais_client.rs | 6 +- .../service-lib/src/service.rs | 58 +++ ...ite_safety_post_annotations_result_sink.py | 1 + .../reply_spam/classifier_coordinated_spam.py | 6 +- grox/flows/reply_spam/task_write.py | 36 +- home-mixer/ads/mod.rs | 2 + home-mixer/ads/multi_risk_blender.rs | 332 ++++++++++++++++++ .../ads/tests/blender_comparison_test.rs | 268 +++++++++++--- home-mixer/ads/tests/mod.rs | 1 + .../ads/tests/multi_risk_blender_tests.rs | 215 ++++++++++++ .../tests/partition_organic_blender_tests.rs | 30 ++ home-mixer/ads/util.rs | 35 ++ .../ads_brand_safety_vf_hydrator.rs | 3 +- .../filters/brazil_2026_election_filter.rs | 140 ++++---- .../filters/viewer_muted_keyword_filter.rs | 18 + home-mixer/params/param.rs | 2 +- home-mixer/selectors/blender_selector.rs | 7 +- home-mixer/server.rs | 13 +- .../response_stats_side_effect.rs | 34 ++ home-mixer/util/mod.rs | 1 + home-mixer/util/strato_context.rs | 75 ++++ phoenix/README.md | 2 +- .../xai-recsys-engine/src/checkpoint_proxy.rs | 4 +- .../xai-recsys-engine/src/copy_port_client.rs | 123 ++++++- .../xai-recsys-engine/src/emb_table.rs | 27 +- .../xai-recsys-engine/src/grpc_compression.rs | 122 ++++++- .../xai-recsys-engine/src/request_metrics.rs | 21 +- .../xai-recsys-proto/proto/recsys.proto | 6 +- .../common/xai-proto/proto/recsys.proto | 6 +- .../xai_checkpointing/load.py | 115 ++++-- phoenix/xrex/configs/xrecsys.py | 1 + phoenix/xrex/train/misc.py | 4 + phoenix/xrex/train/trainer.py | 133 +++++-- .../EntityToSimClustersEmbeddingsJob.scala | 19 +- visibility-filtering/config.rs | 86 ++++- .../safety_label_source/lookup.rs | 84 +++++ .../safety_label_source/metrics.rs | 32 ++ .../safety_label_source/mod.rs | 1 + .../safety_label_source/warmer.rs | 203 +++++++++++ visibility-filtering/server_deps.rs | 60 +++- 41 files changed, 2071 insertions(+), 265 deletions(-) create mode 100644 home-mixer/ads/multi_risk_blender.rs create mode 100644 home-mixer/ads/tests/multi_risk_blender_tests.rs create mode 100644 home-mixer/util/strato_context.rs create mode 100644 visibility-filtering/safety_label_source/warmer.rs diff --git a/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml b/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml index b2ab755c..30a8b0ba 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-25T16:15:48Z +# mirrored from GrowthBook dynamic config; last sync 2026-08-26T16:31:18Z for_entity: user @@ -46,6 +46,8 @@ rules: 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: 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/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/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/task_write.py b/grox/flows/reply_spam/task_write.py index 43d89d13..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,48 +110,45 @@ async def _publish_to_reply_ranking_manhattan( f"Missing user id [_publish_to_reply_ranking_manhattan] {reasoning=} {post.id=} {score=}" ) - existing = None - try: - existing = await ReplyRankingScoreStratoLoader.fetch_reply_ranking_score( - post.id - ) - except Exception: - logger.warning( - f"Failed to fetch existing reply ranking score for {post.id=}, proceeding with write" - ) + existing = await ReplyRankingScoreStratoLoader.fetch_reply_ranking_score( + post.id + ) if ( existing is not None and existing.score is not None - and score > existing.score + and score >= existing.score ): Metrics.counter("task.write_reply_ranking_manhattan.skipped.count").add( - 1, attributes={"reason": "higher_than_existing"} + 1, attributes={"reason": "not_lower_than_existing"} ) logger.info( - f"[_publish_to_reply_ranking_manhattan] skipping write: new {score=} > existing={existing.score} {post.id=}" + 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/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 cb0e9f69..26025ab3 100644 --- a/home-mixer/ads/util.rs +++ b/home-mixer/ads/util.rs @@ -28,6 +28,14 @@ pub(crate) fn has_avoid(post: &ScoredPost) -> bool { ) } +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 +} + pub(crate) fn find_safe_gaps(scored_posts: &[ScoredPost]) -> Vec { let n = scored_posts.len(); let mut safe = Vec::new(); @@ -78,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>, @@ -133,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 9bf776fc..3ac4e132 100644 --- a/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs +++ b/home-mixer/candidate_hydrators/ads_brand_safety_vf_hydrator.rs @@ -144,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}; @@ -480,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 { diff --git a/home-mixer/filters/brazil_2026_election_filter.rs b/home-mixer/filters/brazil_2026_election_filter.rs index 72f685c9..78d92e4b 100644 --- a/home-mixer/filters/brazil_2026_election_filter.rs +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -16,64 +16,70 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // User ids below are obfuscated; usernames are included for transparency. -// @OmarAzizSenador deleted his account at the time this code was written. -// @_ANDREDOPRADO no live account was found. -// @_EDUARDOMANTOAN no live account was found. -// @ADALBERTO_1111 no live account was found. -// @TWITTERADRIANAACCORSI no live account was found. -// @ADRIANASOUSAPIAUI no live account was found. -// @AGoldbach no live account was found. -// @AHELIXO no live account was found. -// @ALCEU_ALCEUMOREIRA no live account was found. -// @ALEXROSETI no live account was found. -// @ALKORAP1 no live account was found. -// @BetoRichaOficial no live account was found. -// @BRUNOPORTODEALMEIDA no live account was found. -// @CHARLES067277 no live account was found. -// @CRISTINAGRAEM no live account was found. -// @DANIELBRSOARES no live account was found. -// @DANILOBALASOFICIAL no live account was found. -// @DANILOTORRES100 no live account was found. -// @DECIOLIMAPT no live account was found. -// @DELEGADOEGUCHI no live account was found. -// @DEMAOLIVEIRA70 no live account was found. -// @DEPCELSOSABINO no live account was found. -// @DEPLUANAREGIA no live account was found. -// @DUARTEJR70 no live account was found. -// @DUDUSIVINSKI no live account was found. -// @EDSONSANTOSRJ no live account was found. -// @EUANGELAGARCIALINKTREE no live account was found. -// @EXPEDITOFUCAP no live account was found. -// @FADAPSICANALISE no live account was found. -// @FEDERALFELICIO no live account was found. -// @GERSONBURMANNIV no live account was found. -// @GSMA1986 no live account was found. -// @JAIZAMETODIO no live account was found. -// @LEOMASCARENHASP no live account was found. -// @LUCIANALIPPI30 no live account was found. -// @LUCIANAOROZIMBO no live account was found. -// @MARCELOSILVACAMPINAS no live account was found. -// @MEUCANA669499 no live account was found. -// @MIRCOCORONETTI no live account was found. -// @NADIAGERHARD no live account was found. -// @NETOFEITOSA6891 no live account was found. -// @PATRICIACRIZANTO2 no live account was found. -// @PAULOMOURAOTO no live account was found. -// @PEDRONASSIF_RJ no live account was found. -// @PEDROPONCIOBE no live account was found. -// @POLICIALPAULOBASTOS no live account was found. -// @SENATORCIDGOMES no live account was found. -// @XIGORPORTO no live account was found. +// 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 @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. +// @ADALBERTO_1111 no live account found. +// @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. +// @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. +// @DECIOLIMAPT no live account found. +// @DELEGADOEGUCHI no live account found. +// @DEMAOLIVEIRA70 no live account found. +// @DEPCELSOSABINO no live account found. +// @DEPLUANAREGIA 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. +// @JAIZAMETODIO no live account found. +// @LEOMASCARENHASP no live account found. +// @LUCIANALIPPI30 no live account found. +// @LUCIANAOROZIMBO 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. +// @OMARAZIZSENADOR no live account found. +// @PATRICIACRIZANTO2 no live account found. +// @PAULOMOURAOTO no live account found. +// @PEDRONASSIF_RJ no live account found. +// @PEDROPONCIOBE no live account found. +// @POLICIALPAULOBASTOS no live account found. +// @PRADOCORONEL no live account found. +// @SENATORCIDGOMES 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([ - // @dayse - 6003262, - // @ricar - 6025402, - // @prado - 9171802, // @madeleinelacsko 9179462, // @renildo @@ -1202,8 +1208,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 129837652, // @alexandrebaldy 130620293, - // @abr - 131428902, // @Miriampetrone 132187525, // @depdelmasso @@ -1244,8 +1248,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 141023529, // @DiogoPBotelho 141087783, - // @PSTUPE - 141090291, // @lucianogenesio 142068227, // @gilmarribeirojr @@ -1462,8 +1464,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 198350574, // @DepEduardoCunha 198535390, - // @acmneto_ - 199025417, // @EdsonSilvaCotia 199102127, // @MariaSeffair @@ -1628,8 +1628,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 276794985, // @danealencar 278126758, - // @paulodimelo - 278319124, // @CostaMarinara 278549268, // @AllanPombopdt @@ -1660,8 +1658,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 289318056, // @waltercamargo40 289521136, - // @ArafetH - 290106695, // @coelho_rodrigo 290204659, // @fefrancischini @@ -1736,12 +1732,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 332324517, // @RequiaoFilho 333720455, - // @cleciacarvalho1 - 334585581, // @Fabinho_Gaspar 334978230, - // @ZanaAmanda - 335065620, // @simboramudar22 337269106, // @_soedi_ @@ -1956,8 +1948,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1081998169, // @deproosevelt 1084884007, - // @AleSilvaOficial - 1089692132, // @depjanetedesa 1094959356, // @NeumannJarbas @@ -1978,8 +1968,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1316758495, // @zecarlospt 1325494376, - // @bispadamares - 1326865753, // @CARLOSVALADARE7 1356677952, // @D_GoretePereira @@ -2764,6 +2752,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1083098735675084800, // @PEDROCO13904182 1083745378489589760, + // @DamaresAlves + 1083774484123975680, // @BennyBriolly 1084272914202091520, // @majorfabianadep @@ -2864,8 +2854,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1109869057153683462, // @DrCesarMello 1110177096373100545, - // @dredsondapaiol - 1110537499242389504, // @ToniettoChris 1110626741314306049, // @keilapereirasp @@ -3334,8 +3322,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1314575606487670784, // @andersonlimaadm 1318194625400737792, - // @SerginhoCaxias - 1318204633823776770, // @delegadalia 1318672895573479424, // @ElianaBayer_ @@ -4911,7 +4897,7 @@ mod tests { #[test] fn hardcoded_list_is_non_empty() { assert!(!BRAZIL_2026_ELECTION_USER_IDS.is_empty()); - assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2328); + assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2315); } #[test] diff --git a/home-mixer/filters/viewer_muted_keyword_filter.rs b/home-mixer/filters/viewer_muted_keyword_filter.rs index ec84c051..482a8182 100644 --- a/home-mixer/filters/viewer_muted_keyword_filter.rs +++ b/home-mixer/filters/viewer_muted_keyword_filter.rs @@ -279,6 +279,24 @@ 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 = ViewerMutedKeywordFilter::new(); diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index db2e30ee..d14c7c94 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-26T16:36:06Z +// mirrored from config feature-switch defaults; last sync 2026-08-27T19:41:17Z use xai_feature_switches::param; param!( 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..3981dcf0 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; @@ -385,6 +386,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { + let polling_header = strato_context::is_polling(request.metadata()); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -392,7 +394,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 = polling_header || proto_query.is_polling; let ctx = self .query_builder .build( @@ -476,6 +478,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { + let polling_header = strato_context::is_polling(request.metadata()); let mut b3_info = extract_b3_info(request.metadata()); b3_info.force_sample(); @@ -486,7 +489,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 = polling_header || proto_query.is_polling; let ctx = self .query_builder .build( @@ -538,6 +541,7 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService &self, request: Request, ) -> Result, Status> { + let polling_header = strato_context::is_polling(request.metadata()); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -545,7 +549,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 = polling_header || proto_query.is_polling; let ctx = self .query_builder .build( @@ -637,6 +641,7 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe &self, request: Request, ) -> Result, Status> { + let polling_header = strato_context::is_polling(request.metadata()); let b3_info = extract_b3_info(request.metadata()); let feed_query = request.into_inner(); let proto_query = feed_query @@ -644,7 +649,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 = polling_header || proto_query.is_polling; let ctx = self .query_builder .build( 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/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/strato_context.rs b/home-mixer/util/strato_context.rs new file mode 100644 index 00000000..44494ca8 --- /dev/null +++ b/home-mixer/util/strato_context.rs @@ -0,0 +1,75 @@ +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)] +struct StratoContext { + #[prost(bool, tag = "11")] + pub is_polling: bool, +} + +pub fn is_polling(metadata: &MetadataMap) -> bool { + extract_strato_context(metadata).is_some_and(|ctx| ctx.is_polling) +} + +fn extract_strato_context(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 + } + + #[test] + fn polling_true() { + assert!(is_polling(&encode_ascii(&StratoContext { + is_polling: true + }))); + } + + #[test] + fn polling_false() { + assert!(!is_polling(&encode_ascii(&StratoContext { + is_polling: false + }))); + } + + #[test] + fn binary_metadata() { + let mut map = MetadataMap::new(); + map.insert_bin( + STRATO_CONTEXT_BIN_KEY, + MetadataValue::from_bytes(&StratoContext { is_polling: true }.encode_to_vec()), + ); + assert!(is_polling(&map)); + } + + #[test] + fn missing_context() { + assert!(!is_polling(&MetadataMap::new())); + } +} diff --git a/phoenix/README.md b/phoenix/README.md index 1d5447c7..f96de574 100644 --- a/phoenix/README.md +++ b/phoenix/README.md @@ -307,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/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs index 6e33ef71..03ba070a 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs @@ -318,9 +318,7 @@ impl CheckpointProxy { let mut handles = Vec::with_capacity(channels.len()); - for (ch_idx, (channel, ch_entries)) in - channels.iter().zip(deduped_entries.into_iter()).enumerate() - { + for (ch_idx, (channel, ch_entries)) in channels.iter().zip(deduped_entries).enumerate() { if ch_entries.is_empty() { continue; } 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 c1e7b5cd..05db78c2 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")] @@ -466,7 +467,7 @@ pub async fn download_dense_and_embeddings( 0.0 }; log::info!( - "copy_port: dense weights loaded prefix={prefix} bytes={bytes} in {:.2}s ({:.2} GB/s)", + "copy_port: dense weights loaded bytes={bytes} in {:.2}s ({:.2} GB/s)", secs, gbs ); @@ -556,10 +557,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) } @@ -607,11 +610,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) } @@ -804,6 +808,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], @@ -891,7 +941,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; @@ -926,6 +976,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; @@ -967,11 +1029,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] = @@ -988,9 +1049,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)] @@ -1134,6 +1194,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!( 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 0673dbbe..68bac048 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs @@ -976,6 +976,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 +1058,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 +1114,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 +1136,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 +1149,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(); 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/request_metrics.rs b/phoenix/crates/serving/xai-recsys-engine/src/request_metrics.rs index 15f1aba7..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); } @@ -706,6 +708,17 @@ mod tests { assert_eq!(src_dc_label(&absent), "unknown"); } + #[test] + fn network_delay_records_under_client_and_src_dc() { + let mut req = Request::new(()); + req.metadata_mut() + .insert("x-request-send-time-ms", "1".parse().unwrap()); + req.metadata_mut().insert("src_dc", "pdxa".parse().unwrap()); + record_client_server_network_delay(&req, "example-service/ranking-xds"); + let _ = + CLIENT_SERVER_NETWORK_DELAY.with_label_values(&["example-service/ranking-xds", "pdxa"]); + } + #[test] fn not_ready_probe_transitions() { disarm_not_ready_probe(); diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index 5d69abef..8228ebc2 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -329,7 +329,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; @@ -1252,6 +1252,8 @@ message SlateContext { optional uint32 sidGap2 = 11; optional uint32 sidGap3 = 12; optional uint32 reconCosMilli = 13; + optional uint32 reconCountAbove = 14; + optional uint32 reconGapAbove = 15; } message ActionInfo { @@ -1601,7 +1603,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; diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index 5d69abef..8228ebc2 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -329,7 +329,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; @@ -1252,6 +1252,8 @@ message SlateContext { optional uint32 sidGap2 = 11; optional uint32 sidGap3 = 12; optional uint32 reconCosMilli = 13; + optional uint32 reconCountAbove = 14; + optional uint32 reconGapAbove = 15; } message ActionInfo { @@ -1601,7 +1603,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; diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py index c10c7a38..33e153f7 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py @@ -12,6 +12,7 @@ import pathlib import re import time +from collections.abc import Iterable from typing import Any, Callable import jax @@ -32,6 +33,54 @@ _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 _prepare_checkpoint_read( + path: pathlib.Path, kms_client: object | None +) -> tuple[bool, object | None, dict[str, Any], list[str], ts.Context]: + encrypted = _is_encrypted_tree(path) + if encrypted: + import xai_kms + + os.environ.setdefault("TENSORSTORE_HTTP_THREADS", "64") + if kms_client is None: + kms_client = xai_kms.KmsClient.from_cluster_env() + metadata_json = json.loads( + xai_kms.nfs.open_envelope(kms_client, str(path / "_METADATA")).read() + ) + checkpoint_names = _names_from_tree_metadata(metadata_json) + else: + 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() + ) + + context_spec = { + "file_io_concurrency": {"limit": 128}, + "cache_pool#ocdbt": {"total_bytes_limit": 100000000}, + } + if encrypted: + context_spec["http_request_concurrency"] = {"limit": 128} + + return encrypted, kms_client, metadata_json, checkpoint_names, ts.Context(context_spec) def _restore_node_serialize_enabled() -> bool: @@ -81,13 +130,13 @@ def _release_batch_memory(): def _build_read_plan( - metadata, + 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 tree_to_dict(metadata, keep_none=False).keys(): + for checkpoint_name in checkpoint_names: name = checkpoint_name if rename is not None: name = rename(checkpoint_name) @@ -153,6 +202,7 @@ def _open_tensor( 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, @@ -162,6 +212,8 @@ def _open_tensor( 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( @@ -254,6 +306,7 @@ def load_checkpoint( 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" @@ -266,18 +319,13 @@ def load_checkpoint( 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}, - } + + encrypted, kms_client, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + path, kms_client ) + use_zarr3 = metadata_json["use_zarr3"] - plan = _build_read_plan(metadata, host_state, load_mask, rename) + 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) @@ -325,7 +373,15 @@ def load_checkpoint( num_batches, ) - try: + with contextlib.ExitStack() as stack: + if node_lock is not None: + stack.callback(node_lock.close) + tspec_transform = None + if encrypted: + import xai_kms + + tspec_transform = stack.enter_context(xai_kms.KvServe(kms_client, path)).rewrite_ocdbt + for batch_index, batch in enumerate(batches, start=1): with node_lock if node_lock is not None else contextlib.nullcontext(): stores = [] @@ -338,6 +394,7 @@ def load_checkpoint( 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)) @@ -355,9 +412,6 @@ def load_checkpoint( _drain_read_futures(futures, host_state, path, timeout, log_loaded=True) del stores _release_batch_memory() - finally: - if node_lock is not None: - node_lock.close() rank_logger.info("Loading checkpoint took %.2f sec", time.time() - start) @@ -409,6 +463,7 @@ def load_checkpoint_streamed( 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" @@ -420,18 +475,12 @@ def load_checkpoint_streamed( start = time.time() 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}, - } + encrypted, kms_client, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + path, kms_client ) + use_zarr3 = metadata_json["use_zarr3"] - plan = _build_read_plan(metadata, device_state, load_mask, rename) + 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 @@ -484,7 +533,15 @@ def load_checkpoint_streamed( num_batches, ) - try: + with contextlib.ExitStack() as stack: + if node_lock is not None: + stack.callback(node_lock.close) + tspec_transform = None + if encrypted: + import xai_kms + + tspec_transform = stack.enter_context(xai_kms.KvServe(kms_client, path)).rewrite_ocdbt + 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}) @@ -499,6 +556,7 @@ def load_checkpoint_streamed( 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)) @@ -524,9 +582,6 @@ def load_checkpoint_streamed( del staging, new_arrays, stores, batch_replaced _release_batch_memory() - finally: - if node_lock is not None: - node_lock.close() rank_logger.info("Loading checkpoint (streamed) took %.2f sec", time.time() - start) return replaced diff --git a/phoenix/xrex/configs/xrecsys.py b/phoenix/xrex/configs/xrecsys.py index 3bd2deb5..34220e3a 100644 --- a/phoenix/xrex/configs/xrecsys.py +++ b/phoenix/xrex/configs/xrecsys.py @@ -641,6 +641,7 @@ def get(self, key: str, default: RecsysTrainer | None = None) -> RecsysTrainer | "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( diff --git a/phoenix/xrex/train/misc.py b/phoenix/xrex/train/misc.py index 87cbeb3b..d0b08db3 100644 --- a/phoenix/xrex/train/misc.py +++ b/phoenix/xrex/train/misc.py @@ -57,6 +57,10 @@ class CheckpointConfig(Config): 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() replication_mode: Literal["full", "dp_only", "none"] = "dp_only" diff --git a/phoenix/xrex/train/trainer.py b/phoenix/xrex/train/trainer.py index ea06718c..512a6b28 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -1177,6 +1177,22 @@ def maybe_load_checkpoint( 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 self.checkpoint_config.save_method != "tensorstore" + ) + 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)) ) @@ -1188,22 +1204,34 @@ def maybe_load_checkpoint( 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 pinned-host staging)" - ) - self.host_state = jax.device_put( - warm_purge(self.state), - warm_purge(restore_staging_sharding), + "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 - self.host_state = jax.device_put(self.state, restore_staging_sharding) + 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) + 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: host_state = self.purge_opt_state_on_load(host_state) @@ -1248,32 +1276,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, - concurrent_gb=self.checkpoint_config.restore_concurrent_gb, - ) + 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, + ) - 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 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 = {} diff --git a/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala b/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala index c1323117..5b9f70b4 100644 --- a/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala +++ b/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala @@ -5,7 +5,10 @@ import com.twitter.recos.entities.thriftscala.Entity import com.twitter.recos.entities.thriftscala.Hashtag import com.twitter.recos.entities.thriftscala.SemanticCoreEntity import com.twitter.scalding._ +import com.twitter.scalding_internal.dalv2.DAL import com.twitter.scalding_internal.dalv2.DALWrite._ +import com.twitter.scalding_internal.dalv2.remote_access.ExplicitLocation +import com.twitter.scalding_internal.dalv2.remote_access.ProcAtla import com.twitter.scalding_internal.multiformat.format.keyval.KeyVal import com.twitter.simclusters_v2.common.ModelVersions import com.twitter.simclusters_v2.common.SimClustersEmbedding @@ -212,8 +215,20 @@ trait EntityToSimClustersEmbeddingApp extends ScheduledExecutionApp { val simClustersEmbedding = jobConfig.modelVersion match { case ModelVersion.Model20m145k2020 => - val simClustersSource2020 = - InterestedInSources.simClustersInterestedIn2020Source(dateRange, timeZone) + val interestedIn2020WithFallback = + SimclustersV2InterestedIn20M145K2020ScalaDataset.copy(fallbackPath = Some("viewfs://hadoop-nn.example.invalid/user/cassowary/manhattan_sequence_files/" + + s"simclusters_v2_interested_in_20M_145K_2020/_TMP_RECOVERY/${dateRange.start.timestamp}")) + val simClustersSource2020 = DAL + .readMostRecentSnapshot( + interestedIn2020WithFallback, + dateRange.prepend(Days(28)(timeZone)) + ) + .withRemoteReadPolicy(ExplicitLocation(ProcAtla)) + .toTypedPipe + .map { + case KeyVal(userId, clustersUserIsInterestedIn) => + (userId, clustersUserIsInterestedIn) + } computeEmbeddings( simClustersSource2020, normalizedUserEntityMatrix, diff --git a/visibility-filtering/config.rs b/visibility-filtering/config.rs index 1eff08b3..c67f1ea1 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -6,6 +6,34 @@ pub const ENV_GRPC_MTLS_CLIENT_CA_PATH: &str = "GRPC_MTLS_CLIENT_CA_PATH"; pub const ENV_DUAL_CALL_HARNESS_ENABLED: &str = "VF_DUAL_CALL_HARNESS_ENABLED"; 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_CACHE_WARM_SAMPLE_PCT: &str = "VF_CACHE_WARM_SAMPLE_PCT"; +pub const ENV_APP_ENV: &str = "APP_ENV"; +pub const ENV_GIZMODUCK_CLIENT_ID: &str = "VF_GIZMODUCK_CLIENT_ID"; +pub const ENV_TWEMCACHE_CLIENT_NAME: &str = "VF_TWEMCACHE_CLIENT_NAME"; + +pub 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 twemcache_client_name() -> String { + resolve_twemcache_client_name(std::env::var(ENV_TWEMCACHE_CLIENT_NAME).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 resolve_twemcache_client_name(configured: Option<&str>) -> String { + configured + .unwrap_or("visibility-filtering-service") + .to_string() +} pub fn dual_call_harness_enabled() -> bool { parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) @@ -27,6 +55,21 @@ pub fn fallback_cache_populate_enabled() -> bool { ) } +pub fn cache_warm_sample_pct() -> u8 { + parse_sample_pct(std::env::var(ENV_CACHE_WARM_SAMPLE_PCT).ok().as_deref()) + .unwrap_or_else(|error| panic!("{ENV_CACHE_WARM_SAMPLE_PCT}: {error}")) +} + +fn parse_sample_pct(value: Option<&str>) -> Result { + let Some(value) = value.map(str::trim).filter(|v| !v.is_empty()) else { + return Ok(0); + }; + match value.parse::() { + Ok(pct) if pct <= 100 => Ok(pct), + _ => Err(format!("expected an integer 0-100, got {value:?}")), + } +} + fn parse_env_flag(value: Option<&str>) -> bool { value.is_some_and(|value| { matches!( @@ -103,7 +146,48 @@ impl GrpcMtlsConfig { #[cfg(test)] mod tests { - use super::parse_env_flag; + use super::{ + parse_env_flag, parse_sample_pct, resolve_gizmoduck_client_id, + resolve_twemcache_client_name, + }; + + #[test] + fn parses_sample_pct_range() { + assert_eq!(parse_sample_pct(None), Ok(0)); + assert_eq!(parse_sample_pct(Some("")), Ok(0)); + assert_eq!(parse_sample_pct(Some("100")), Ok(100)); + } + + #[test] + fn rejects_invalid_sample_pct() { + for value in ["101", "on"] { + assert!(parse_sample_pct(Some(value)).is_err(), "{value}"); + } + } + + #[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/safety_label_source/lookup.rs b/visibility-filtering/safety_label_source/lookup.rs index 36e8fd04..b1c4163a 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::SampledWarmer; + + let (warmer, _rx) = SampledWarmer::without_drain_task(1, 100); + 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/metrics.rs b/visibility-filtering/safety_label_source/metrics.rs index 770eb6af..695b65db 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,37 @@ pub(crate) fn record_cache_fallback_keys( ); } +#[derive(Clone, Copy)] +pub(crate) enum WarmKeyResult { + EligibleMiss, + SampledOut, + Enqueued, + DroppedChannelFull, + FetchIssued, + FetchFailed, +} + +impl WarmKeyResult { + fn as_str(self) -> &'static str { + match self { + Self::EligibleMiss => "eligible_miss", + Self::SampledOut => "sampled_out", + 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..008d404f 100644 --- a/visibility-filtering/safety_label_source/mod.rs +++ b/visibility-filtering/safety_label_source/mod.rs @@ -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/warmer.rs b/visibility-filtering/safety_label_source/warmer.rs new file mode 100644 index 00000000..18146efa --- /dev/null +++ b/visibility-filtering/safety_label_source/warmer.rs @@ -0,0 +1,203 @@ +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 SampledWarmer { + tx: mpsc::Sender>, + sample_pct: u8, +} + +impl SampledWarmer { + pub(crate) fn spawn(fetcher: Arc, sample_pct: u8) -> Arc { + let (tx, rx) = mpsc::channel(WARM_CHANNEL_CAPACITY); + tokio::spawn(drain(rx, fetcher)); + Arc::new(Self { tx, sample_pct }) + } + + #[cfg(test)] + pub(crate) fn without_drain_task( + capacity: usize, + sample_pct: u8, + ) -> (Self, mpsc::Receiver>) { + let (tx, rx) = mpsc::channel(capacity); + (Self { tx, sample_pct }, rx) + } +} + +impl Warmer for SampledWarmer { + fn warm(&self, mut miss_ids: Vec) { + metrics::record_cache_warm_keys(WarmKeyResult::EligibleMiss, miss_ids.len()); + if self.sample_pct < 100 { + let eligible = miss_ids.len(); + miss_ids.retain(|_| fastrand::u8(..100) < self.sample_pct); + metrics::record_cache_warm_keys(WarmKeyResult::SampledOut, eligible - 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 = SampledWarmer::spawn(fetcher.clone(), 100); + + 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 = SampledWarmer::spawn(fetcher.clone(), 100); + + 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 = SampledWarmer::spawn(fetcher.clone(), 100); + + 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) = SampledWarmer::without_drain_task(1, 100); + + 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_deps.rs b/visibility-filtering/server_deps.rs index 2da74a1c..99b4fe3b 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -9,6 +9,7 @@ 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::{SampledWarmer, StratoWarmFetcher, Warmer}; use crate::safety_label_source::{ManhattanLabelFetcher, MhLabelClient, SafetyLabelSource}; use crate::server::VFServer; use std::future::Future; @@ -114,10 +115,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( @@ -165,11 +163,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,9 +189,15 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { 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( @@ -260,6 +265,49 @@ async fn build_reference_compare_harness( Some(Arc::new(ReferenceCompareHarness::new(strato, datacenter))) } +const CACHE_WARM_REQUEST_TIMEOUT_MS: u64 = 500; + +async fn build_cache_warmer( + datacenter: &str, + init_deadline: tokio::time::Instant, +) -> Option> { + let sample_pct = crate::config::cache_warm_sample_pct(); + if sample_pct == 0 { + 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!(sample_pct, "L2 cache warmer enabled"); + Some(SampledWarmer::spawn( + Arc::new(StratoWarmFetcher::new(grpc)), + sample_pct, + )) +} + const TES_STRATO_REQUEST_TIMEOUT_MS: u64 = 100; fn tes_client_config(deterministic_aperture: bool) -> TESClientConfig { From bc8e5f0f07b31337bfdcaf690121498e00199b64 Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 28 Aug 2026 22:37:32 +0000 Subject: [PATCH 11/18] Open-source X Recommendation Algorithm --- README.md | 2 +- home-mixer/models/candidate.rs | 3 + home-mixer/params/param.rs | 3 +- home-mixer/scorers/phoenix_scorer.rs | 4 +- home-mixer/util/phoenix_request.rs | 5 +- .../xai-recsys-engine/src/copy_port_client.rs | 31 +- .../xai-recsys-engine/src/emb_table.rs | 91 +- .../xai-recsys-engine/src/proto_parser.rs | 11 + .../xai-recsys-proto/proto/recsys.proto | 3 + .../common/xai-proto/proto/recsys.proto | 3 + visibility-filtering/rules/context.rs | 214 ++++ visibility-filtering/rules/fixtures.rs | 116 ++ visibility-filtering/rules/golden_corpus.rs | 1039 +++++++++++++++++ visibility-filtering/rules/mod.rs | 61 +- visibility-filtering/rules/nsfw_age_gating.rs | 175 ++- .../rules/nsfw_interstitial.rs | 109 +- visibility-filtering/rules/nullcast_rule.rs | 71 +- visibility-filtering/rules/registry.rs | 347 +----- .../rules/socialgraph_rules.rs | 126 +- visibility-filtering/rules/tes_rules.rs | 348 ++---- .../rules/tweet_flag_rules.rs | 47 +- .../rules/tweet_label_drops.rs | 131 +-- .../rules/user_label_drops.rs | 72 +- visibility-filtering/rules/user_rules.rs | 113 +- 24 files changed, 1950 insertions(+), 1175 deletions(-) create mode 100644 visibility-filtering/rules/context.rs create mode 100644 visibility-filtering/rules/fixtures.rs create mode 100644 visibility-filtering/rules/golden_corpus.rs diff --git a/README.md b/README.md index e4c5656e..c70eba36 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ This repository contains the core code that determines which posts a viewer sees 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 25, 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). +- **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 diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index 49d5fa88..ddccebc3 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -24,6 +24,8 @@ pub struct PostCandidate { pub slate_context: Option, #[serde(default)] pub served_slate_context: Option, + #[serde(default)] + pub reranker_head_tag: Option, #[serde( serialize_with = "serialize_served_type", deserialize_with = "deserialize_served_type" @@ -208,6 +210,7 @@ impl CandidateHelpers for PostCandidate { recon_gap_above: c.recon_gap_above, }), reward_rerank_slot_prob: None, + reranker_head_tag: self.reranker_head_tag, } } diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index d14c7c94..80917b59 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-27T19:41:17Z +// mirrored from config feature-switch defaults; last sync 2026-08-28T20:07:44Z use xai_feature_switches::param; param!( @@ -243,6 +243,7 @@ param!( "rust_home_mixer_log_slate_context", false ); +param!(RerankerHeadTag, i64, "rust_home_mixer_reranker_head_tag", 0); param!( OonWeightFactor, f64, diff --git a/home-mixer/scorers/phoenix_scorer.rs b/home-mixer/scorers/phoenix_scorer.rs index 5cf292d7..bba14c3d 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; @@ -112,6 +112,7 @@ impl Scorer for PhoenixScorer { .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) @@ -123,5 +124,6 @@ impl Scorer for PhoenixScorer { 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/util/phoenix_request.rs b/home-mixer/util/phoenix_request.rs index 96635b5a..ab58a3e4 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::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; @@ -145,7 +146,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/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/copy_port_client.rs index 05db78c2..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 @@ -234,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, @@ -249,11 +259,7 @@ async fn run_downloads( .max(1); join_rate_limited(futures, limit, max_c).await } - _ => 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}"))), + _ => join_transfers(futures).await, } } @@ -273,11 +279,7 @@ async fn join_rate_limited( break; } let batch_size = batch.len(); - let batch_results = join_all(batch.into_iter().map(tokio::task::spawn)) - .await - .into_iter() - .collect::, _>>() - .map_err(|e| CopyPortError::Other(format!("copy_port download task join: {e}")))?; + let batch_results = join_transfers(batch).await?; let failed = batch_results .iter() .filter(|r| r.0 == TRANSFER_FAILED_SENTINEL) @@ -1351,7 +1353,7 @@ mod tests { } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn rate_limit_caps_in_flight_despite_spawn() { + 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])); @@ -1366,14 +1368,13 @@ mod tests { 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], - "join_all on JoinHandles must keep submission order" + 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_spawned_batches() { + 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])); @@ -1405,7 +1406,7 @@ mod tests { let between = second_batch_start.saturating_duration_since(first_batch_start); assert!( between >= Duration::from_millis(700), - "second batch started {between:?} after the first; expected ~1s pacing sleep" + "second batch started {between:?} after the first" ); let total_bytes = (4 * bytes) as f64; 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 68bac048..5a561b50 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs @@ -610,6 +610,45 @@ pub(crate) async fn expand_replicated_channels( ); } +struct JoinOnDrop { + 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 +675,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 +709,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 +727,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) => { 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-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index 8228ebc2..e9100296 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -1236,6 +1236,7 @@ message ScoreInfo { optional double finalScore = 3; SlateContext slateContext = 4; optional double rewardRerankSlotProb = 5; + optional uint32 rerankerHeadTag = 6; } message SlateContext { @@ -1311,6 +1312,8 @@ message UserActionSequence { AdUserInfo adUserInfo = 7; repeated Mask masks = 8; + + optional uint32 rerankerHeadTag = 9; } message UserActionList { diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index 8228ebc2..e9100296 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -1236,6 +1236,7 @@ message ScoreInfo { optional double finalScore = 3; SlateContext slateContext = 4; optional double rewardRerankSlotProb = 5; + optional uint32 rerankerHeadTag = 6; } message SlateContext { @@ -1311,6 +1312,8 @@ message UserActionSequence { AdUserInfo adUserInfo = 7; repeated Mask masks = 8; + + optional uint32 rerankerHeadTag = 9; } message UserActionList { diff --git a/visibility-filtering/rules/context.rs b/visibility-filtering/rules/context.rs new file mode 100644 index 00000000..5a5a0135 --- /dev/null +++ b/visibility-filtering/rules/context.rs @@ -0,0 +1,214 @@ +use crate::models::{HydratedTweetCandidate, SafetyLabelType, ViewerFeatures}; +use crate::rules::registry::SafetyLevel; +use xai_core_entities::entities::TakedownReason; +use xai_x_thrift::user_labels::LabelValue; + +pub struct RuleContext<'a> { + safety_level: SafetyLevel, + viewer: &'a ViewerFeatures, + candidate: &'a HydratedTweetCandidate, +} + +impl<'a> RuleContext<'a> { + pub(super) 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 viewer_is_logged_out(&self) -> bool { + self.viewer.viewer_is_logged_out() + } + + pub fn viewer_is_underage(&self) -> bool { + self.viewer.viewer_is_underage() + } + + pub fn viewer_has_no_stated_age(&self) -> bool { + self.viewer.viewer_has_no_stated_age() + } + + pub fn viewer_allows_sensitive_media(&self) -> bool { + self.viewer.allows_sensitive_media + } + + pub fn viewer_country_in(&self, countries: &[&str]) -> bool { + self.viewer + .account_country_code + .as_deref() + .or(self.viewer.country_code.as_deref()) + .is_some_and(|c| countries.contains(&c)) + } + + 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_blocks_author(&self) -> bool { + self.candidate.relationship.viewer_blocks_author + } + + pub fn viewer_mutes_author(&self) -> bool { + self.candidate.relationship.viewer_mutes_author + } + + pub fn viewer_mutes_retweets_from_author(&self) -> bool { + self.candidate + .relationship + .viewer_mutes_retweets_from_author + } + + pub fn has_tweet_safety_label(&self, label: SafetyLabelType) -> bool { + self.candidate.has_safety_label(label) + } + + pub fn is_retweet(&self) -> bool { + self.candidate.is_retweet() + } + + pub fn is_stale_tweet(&self) -> bool { + self.candidate.is_stale_tweet() + } + + pub fn is_nullcast(&self) -> bool { + self.candidate.is_nullcast() + } + + pub fn is_community_tweet(&self) -> bool { + self.candidate.is_community_tweet() + } + + pub fn has_media(&self) -> bool { + self.candidate.has_media() + } + + pub fn has_dmca_media(&self) -> bool { + self.candidate.has_dmca_media() + } + + pub fn is_nsfw_flagged(&self) -> bool { + self.candidate.is_nsfw_flagged() + } + + pub fn has_tweet_nsfw_user_flag(&self) -> bool { + self.candidate.tweet_features.nsfw.user + } + + pub fn has_tweet_nsfw_admin_flag(&self) -> bool { + self.candidate.tweet_features.nsfw.admin + } + + pub fn legal_takedown_in_viewer_country(&self) -> bool { + self.takedown_in_viewer_country(legal_takedown_country) + } + + pub fn local_laws_takedown_in_viewer_country(&self) -> bool { + self.takedown_in_viewer_country(local_laws_takedown_country) + } + + fn takedown_in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool { + let Some(viewer_country) = &self.viewer.country_code else { + return false; + }; + self.candidate + .tweet_features + .takedown + .reasons + .iter() + .filter_map(extractor) + .any(|c| c.eq_ignore_ascii_case(viewer_country)) + } + + pub fn media_restricted_in_viewer_country(&self) -> bool { + let country = self + .viewer + .country_code + .as_deref() + .unwrap_or(WORLDWIDE_COUNTRY_CODE); + let allow = &self.candidate.tweet_features.media.geo_allow_list; + let deny = &self.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)) + } + + pub fn author_is_suspended(&self) -> bool { + self.candidate.author_features.is_suspended + } + + pub fn author_is_deactivated(&self) -> bool { + self.candidate.author_features.is_deactivated + } + + pub fn author_is_erased(&self) -> bool { + self.candidate.author_features.is_erased + } + + pub fn author_is_offboarded(&self) -> bool { + self.candidate.author_features.is_offboarded + } + + pub fn author_is_protected(&self) -> bool { + self.candidate.author_features.is_protected + } + + pub fn author_is_nsfw_user(&self) -> bool { + self.candidate.author_features.is_nsfw_user + } + + pub fn author_is_nsfw_admin(&self) -> bool { + self.candidate.author_features.is_nsfw_admin + } + + pub fn author_has_user_label(&self, label: LabelValue) -> bool { + self.candidate.author_has_user_label(label) + } + + pub fn is_exclusive_tweet(&self) -> bool { + self.candidate.exclusive_content.is_some() + } + + pub fn viewer_is_conversation_author(&self) -> bool { + match (&self.candidate.exclusive_content, self.viewer.viewer_id()) { + (Some(exclusive), Some(viewer_id)) => viewer_id == exclusive.conversation_author_id, + _ => false, + } + } + + pub fn viewer_super_follows_author(&self) -> bool { + self.candidate + .exclusive_content + .as_ref() + .is_some_and(|exclusive| exclusive.viewer_super_follows_author) + } +} + +const WORLDWIDE_COUNTRY_CODE: &str = "xx"; + +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, + } +} diff --git a/visibility-filtering/rules/fixtures.rs b/visibility-filtering/rules/fixtures.rs new file mode 100644 index 00000000..b0c491ed --- /dev/null +++ b/visibility-filtering/rules/fixtures.rs @@ -0,0 +1,116 @@ +use crate::models::{ + AuthorFeatures, HydratedTweetCandidate, SafetyLabel, SafetyLabelMap, SafetyLabelType, + TweetFeatures, UserLabelSet, Viewer, ViewerAuthorRelationship, ViewerFeatures, +}; +use std::collections::{HashMap, HashSet}; +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(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: HashMap::new(), + user_labels: HashSet::new(), + } +} + +pub(crate) struct CandidateBuilder { + candidate: HydratedTweetCandidate, + labels: HashMap, + 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, SafetyLabel::default()); + 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 + } +} diff --git a/visibility-filtering/rules/golden_corpus.rs b/visibility-filtering/rules/golden_corpus.rs new file mode 100644 index 00000000..dbc3f31c --- /dev/null +++ b/visibility-filtering/rules/golden_corpus.rs @@ -0,0 +1,1039 @@ +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::{Policies, 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 policies = Policies::new(); + let mut failures = Vec::new(); + for case in cases() { + let verdict = policies.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 policies = Policies::new(); + let wired: BTreeSet<&'static str> = [FilterAll, TimelineHome, TimelineHomeRecommendations] + .into_iter() + .flat_map(|level| policies.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 Policies 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 { + let mut features = TweetFeatures::default(); + features.takedown.reasons = vec![reason]; + candidate().with_tweet_features(features).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, + }, + ] +} + +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"), + }, + ] +} + +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, + }, + ] +} + +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/mod.rs b/visibility-filtering/rules/mod.rs index 2750312a..33f7df2a 100644 --- a/visibility-filtering/rules/mod.rs +++ b/visibility-filtering/rules/mod.rs @@ -1,3 +1,8 @@ +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; @@ -10,59 +15,12 @@ pub mod tweet_label_drops; pub mod user_label_drops; pub mod user_rules; -use crate::models::{HydratedTweetCandidate, SafetyLabelType, VfAction, ViewerFeatures}; +use crate::models::VfAction; use xai_visibility_filtering::models::FilteredReason; +pub use context::RuleContext; 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; @@ -111,8 +69,8 @@ fn evaluate_rules(rules: &[Box], context: &RuleContext<'_>) -> Verdict #[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) } @@ -120,6 +78,7 @@ pub(crate) fn test_context<'a>( #[cfg(test)] mod tests { use super::*; + use crate::models::{HydratedTweetCandidate, ViewerFeatures}; use std::sync::{Arc, Mutex}; struct FakeRule { diff --git a/visibility-filtering/rules/nsfw_age_gating.rs b/visibility-filtering/rules/nsfw_age_gating.rs index a2a0e65e..1146c2b2 100644 --- a/visibility-filtering/rules/nsfw_age_gating.rs +++ b/visibility-filtering/rules/nsfw_age_gating.rs @@ -1,4 +1,4 @@ -use crate::models::{SafetyLabelType, VfAction, ViewerFeatures}; +use crate::models::{SafetyLabelType, VfAction}; use crate::rules::{Rule, RuleContext}; use xai_visibility_filtering::models::FilteredReason; @@ -7,18 +7,16 @@ const NSFW_GATING_COUNTRIES: [&str; 16] = [ ]; fn nsfw_base_condition(context: &RuleContext<'_>) -> bool { - let candidate = context.candidate(); !context.is_author_viewer() - && candidate.has_media() + && context.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())) + || (context.is_nsfw_flagged() && !context.is_retweet())) } fn graphic_base_condition(context: &RuleContext<'_>) -> bool { - let candidate = context.candidate(); !context.is_author_viewer() - && candidate.has_media() + && context.has_media() && context.has_tweet_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION) } @@ -34,14 +32,6 @@ fn sensitive_base_condition(context: &RuleContext<'_>) -> bool { || 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 { @@ -50,8 +40,7 @@ impl Rule for SensitiveViewerLoggedOutDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_is_logged_out() && sensitive_base_condition(context) { + if context.viewer_is_logged_out() && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) } else { VfAction::Allow @@ -67,8 +56,7 @@ impl Rule for SensitiveViewerUnderageDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_is_underage() && sensitive_base_condition(context) { + if context.viewer_is_underage() && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) } else { VfAction::Allow @@ -84,9 +72,8 @@ impl Rule for SensitiveViewerNoStatedAgeDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - if viewer.viewer_has_no_stated_age() - && viewer_in_gating_jurisdiction(viewer) + if context.viewer_has_no_stated_age() + && context.viewer_country_in(&NSFW_GATING_COUNTRIES) && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) @@ -100,56 +87,30 @@ impl Rule for SensitiveViewerNoStatedAgeDropRule { mod tests { use super::*; use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, MediaFeature, NsfwFeature, SafetyLabel, - SafetyLabelMap, TweetFeatures, Viewer, ViewerAge, ViewerFeatures, + AuthorFeatures, HydratedTweetCandidate, NsfwFeature, Viewer, ViewerAge, ViewerFeatures, }; - use std::collections::HashMap; + use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; - fn viewer(age: ViewerAge) -> ViewerFeatures { + fn gating_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, + ..viewer(VIEWER_ID) } } 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() - } + candidate().with_label(label).with_media().build() } 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 { + candidate() + .with_media() + .with_author_features(AuthorFeatures { is_nsfw_user: true, ..Default::default() - }, - ..Default::default() - } + }) + .build() } #[test] @@ -157,7 +118,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -169,7 +130,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_RECALL); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -181,7 +142,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -193,7 +154,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); let v = ViewerFeatures { allows_sensitive_media: true, - ..viewer(ViewerAge::Known(15)) + ..gating_viewer(ViewerAge::Known(15)) }; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -206,7 +167,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(18)), + &gating_viewer(ViewerAge::Known(18)), &c )), VfAction::Allow @@ -217,13 +178,17 @@ mod tests { 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)), + SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( + &gating_viewer(ViewerAge::Unknown), + &c + )), VfAction::Allow )); assert!(matches!( - SensitiveViewerNoStatedAgeDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), + SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( + &gating_viewer(ViewerAge::Unknown), + &c + )), VfAction::Allow )); } @@ -239,7 +204,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -251,7 +216,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_CARD_IMAGE); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -263,7 +228,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::NotStated), + &gating_viewer(ViewerAge::NotStated), &c )), VfAction::Drop(_) @@ -275,7 +240,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); let v = ViewerFeatures { country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -288,7 +253,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); let v = ViewerFeatures { viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) + ..gating_viewer(ViewerAge::Unknown) }; assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -301,7 +266,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_CARD_IMAGE); let v = ViewerFeatures { viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) + ..gating_viewer(ViewerAge::Unknown) }; assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -314,7 +279,7 @@ mod tests { let c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(18)), + &gating_viewer(ViewerAge::Known(18)), &c )), VfAction::Allow @@ -325,13 +290,17 @@ mod tests { 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)), + SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( + &gating_viewer(ViewerAge::Unknown), + &c + )), VfAction::Allow )); assert!(matches!( - SensitiveViewerNoStatedAgeDropRule - .evaluate(&crate::rules::test_context(&viewer(ViewerAge::Unknown), &c)), + SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( + &gating_viewer(ViewerAge::Unknown), + &c + )), VfAction::Allow )); } @@ -339,10 +308,10 @@ mod tests { #[test] fn nsfw_text_self_view_is_exempt() { let mut c = no_media_candidate_with_label(SafetyLabelType::NSFW_TEXT); - c.author_id = 999; + c.author_id = VIEWER_ID; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -355,7 +324,7 @@ mod tests { c.tweet_features.media.has_media = false; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -365,10 +334,10 @@ mod tests { #[test] fn self_view_is_exempt() { let mut c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.author_id = 999; + c.author_id = VIEWER_ID; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -380,7 +349,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::NotStated), + &gating_viewer(ViewerAge::NotStated), &c )), VfAction::Drop(_) @@ -392,7 +361,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); let v = ViewerFeatures { country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -405,7 +374,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); let v = ViewerFeatures { country_code: None, - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -419,7 +388,7 @@ mod tests { let v = ViewerFeatures { country_code: Some("de".into()), account_country_code: Some("us".into()), - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -433,7 +402,7 @@ mod tests { let v = ViewerFeatures { country_code: Some("us".into()), account_country_code: Some("kr".into()), - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -447,7 +416,7 @@ mod tests { let v = ViewerFeatures { country_code: Some("de".into()), account_country_code: None, - ..viewer(ViewerAge::NotStated) + ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -460,7 +429,7 @@ mod tests { let c = nsfw_author_media_candidate(); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -477,7 +446,7 @@ mod tests { }; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -499,7 +468,7 @@ mod tests { let c = nsfw_tweet_flag_media_candidate(); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -515,7 +484,7 @@ mod tests { }; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -528,7 +497,7 @@ mod tests { c.author_features.is_nsfw_user = true; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Drop(_) @@ -541,7 +510,7 @@ mod tests { c.tweet_features.nsfw = NsfwFeature::default(); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -554,7 +523,7 @@ mod tests { c.tweet_features.core.source_tweet_id = Some(42); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -564,10 +533,10 @@ mod tests { #[test] fn nsfw_tweet_flag_self_view_exempt() { let mut c = nsfw_tweet_flag_media_candidate(); - c.author_id = 999; + c.author_id = VIEWER_ID; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -580,7 +549,7 @@ mod tests { c.tweet_features.core.source_tweet_id = Some(42); assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -593,7 +562,7 @@ mod tests { c.tweet_features.media.has_media = false; assert!(matches!( SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow @@ -605,7 +574,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); let v = ViewerFeatures { viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) + ..gating_viewer(ViewerAge::Unknown) }; assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -618,7 +587,7 @@ mod tests { let c = nsfw_author_media_candidate(); let v = ViewerFeatures { viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) + ..gating_viewer(ViewerAge::Unknown) }; assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -632,7 +601,7 @@ mod tests { c.tweet_features.media.has_media = false; let v = ViewerFeatures { viewer: Viewer::LoggedOut, - ..viewer(ViewerAge::Unknown) + ..gating_viewer(ViewerAge::Unknown) }; assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context(&v, &c)), @@ -645,7 +614,7 @@ mod tests { let c = media_candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); assert!(matches!( SensitiveViewerLoggedOutDropRule.evaluate(&crate::rules::test_context( - &viewer(ViewerAge::Known(15)), + &gating_viewer(ViewerAge::Known(15)), &c )), VfAction::Allow diff --git a/visibility-filtering/rules/nsfw_interstitial.rs b/visibility-filtering/rules/nsfw_interstitial.rs index c1d25ee3..674de6fa 100644 --- a/visibility-filtering/rules/nsfw_interstitial.rs +++ b/visibility-filtering/rules/nsfw_interstitial.rs @@ -55,9 +55,8 @@ impl Rule for NsfwAuthorInterstitialRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.is_nsfw_flagged() - && candidate.has_media() + if context.is_nsfw_flagged() + && context.has_media() && !context.is_author_viewer() && !context.viewer_allows_sensitive_media() { @@ -70,107 +69,72 @@ impl Rule for NsfwAuthorInterstitialRule { #[cfg(test)] mod tests { use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, MediaFeature, NsfwFeature, SafetyLabel, - SafetyLabelMap, TweetFeatures, Viewer, ViewerFeatures, + use crate::models::{AuthorFeatures, HydratedTweetCandidate, NsfwFeature}; + use crate::rules::fixtures::{ + author_viewer, candidate, sensitive_opt_in_viewer, viewer, VIEWER_ID, }; - 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); + let c = candidate() + .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) + .build(); assert!(matches!( NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Interstitial(_) )); } #[test] fn interstitial_allows_opt_in() { - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); + let c = candidate() + .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) + .build(); assert!(matches!( NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_sensitive_opt_in(), &c)), + .evaluate(&crate::rules::test_context(&sensitive_opt_in_viewer(), &c)), VfAction::Allow )); } #[test] fn interstitial_allows_self_view() { - let mut c = candidate_with_label(SafetyLabelType::NSFW_HIGH_PRECISION); - c.author_id = 999; + let c = candidate() + .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) + .build(); assert!(matches!( NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), + .evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } #[test] fn interstitial_allows_no_label() { - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; + let c = candidate().build(); assert!(matches!( NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer_default(), &c)), + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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 { + candidate() + .with_media() + .with_author_features(AuthorFeatures { is_nsfw_user: true, ..Default::default() - }, - ..Default::default() - } + }) + .build() } #[test] fn author_interstitial_blurs_non_opt_in() { assert!(matches!( NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_default(), + &viewer(VIEWER_ID), &nsfw_author_candidate() )), VfAction::Interstitial(_) @@ -181,7 +145,7 @@ mod tests { fn author_interstitial_allows_opt_in() { assert!(matches!( NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_sensitive_opt_in(), + &sensitive_opt_in_viewer(), &nsfw_author_candidate() )), VfAction::Allow @@ -193,7 +157,8 @@ mod tests { let mut c = nsfw_author_candidate(); c.tweet_features.media.has_media = false; assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), + NsfwAuthorInterstitialRule + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -212,7 +177,7 @@ mod tests { fn tweet_flag_interstitial_blurs_non_opt_in() { assert!(matches!( NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_default(), + &viewer(VIEWER_ID), &nsfw_tweet_flag_candidate() )), VfAction::Interstitial(_) @@ -227,7 +192,8 @@ mod tests { admin: true, }; assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), + NsfwAuthorInterstitialRule + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Interstitial(_) )); } @@ -237,7 +203,8 @@ mod tests { 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)), + NsfwAuthorInterstitialRule + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Interstitial(_) )); } @@ -247,7 +214,8 @@ mod tests { let mut c = nsfw_tweet_flag_candidate(); c.tweet_features.nsfw = NsfwFeature::default(); assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), + NsfwAuthorInterstitialRule + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -256,7 +224,7 @@ mod tests { fn tweet_flag_interstitial_allows_opt_in() { assert!(matches!( NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer_sensitive_opt_in(), + &sensitive_opt_in_viewer(), &nsfw_tweet_flag_candidate() )), VfAction::Allow @@ -265,10 +233,9 @@ mod tests { #[test] fn tweet_flag_interstitial_allows_self_view() { - let mut c = nsfw_tweet_flag_candidate(); - c.author_id = 999; + let c = nsfw_tweet_flag_candidate(); assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&viewer_default(), &c)), + NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } diff --git a/visibility-filtering/rules/nullcast_rule.rs b/visibility-filtering/rules/nullcast_rule.rs index d6ce983f..db4e22b7 100644 --- a/visibility-filtering/rules/nullcast_rule.rs +++ b/visibility-filtering/rules/nullcast_rule.rs @@ -10,8 +10,7 @@ impl Rule for NullcastedTweetDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.is_nullcast() && !candidate.is_retweet() && !candidate.is_community_tweet() { + if context.is_nullcast() && !context.is_retweet() && !context.is_community_tweet() { return VfAction::Drop(FilteredReason::TweetIsNullcast); } VfAction::Allow @@ -21,30 +20,24 @@ impl Rule for NullcastedTweetDropRule { #[cfg(test)] mod tests { use super::*; - use crate::models::{ - CoreFeature, HydratedTweetCandidate, TweetFeatures, Viewer, ViewerFeatures, - }; + use crate::models::{HydratedTweetCandidate, TweetFeatures}; + use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } + fn nullcast_candidate() -> HydratedTweetCandidate { + candidate() + .with_tweet_features(TweetFeatures { + is_nullcast: true, + ..Default::default() + }) + .build() } #[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() - }; + let c = nullcast_candidate(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -52,17 +45,10 @@ mod tests { #[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() - }; + let mut c = nullcast_candidate(); + c.tweet_features.is_community_tweet = true; assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -70,20 +56,10 @@ mod tests { #[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() - }; + let mut c = nullcast_candidate(); + c.tweet_features.core.source_tweet_id = Some(99); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -91,16 +67,9 @@ mod tests { #[test] fn non_nullcast_allows() { let rule = NullcastedTweetDropRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { - is_nullcast: false, - ..Default::default() - }, - ..Default::default() - }; + let c = candidate().build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index 64ac414a..acf6f96f 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -72,6 +72,11 @@ impl Policies { evaluate_rules(self.select(level), &context) } + #[cfg(test)] + pub(crate) fn wired_rule_names(&self, level: SafetyLevel) -> Vec<&'static str> { + self.select(level).iter().map(|rule| rule.name()).collect() + } + pub fn rule_counts(&self) -> (usize, usize) { ( self.timeline_home.len(), @@ -172,10 +177,8 @@ fn timeline_home_recommendations_policy() -> Vec> { #[cfg(test)] mod tests { use super::*; - use crate::models::{ - HydratedTweetCandidate, MediaFeature, TweetFeatures, Viewer, ViewerFeatures, - }; - use std::collections::HashMap; + use crate::models::{HydratedTweetCandidate, MediaFeature, TweetFeatures, ViewerFeatures}; + use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; struct RecommendationsOnlyRule; @@ -224,15 +227,8 @@ mod tests { #[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() - }; + let candidate = candidate().build(); + let viewer = author_viewer(); assert!(matches!( FilterAllRule.evaluate(&crate::rules::test_context(&viewer, &candidate)), VfAction::Drop(_) @@ -242,11 +238,7 @@ mod tests { #[test] fn filter_all_policy_drops_pristine_candidate() { let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; + let candidate = candidate().build(); let verdict = policies.evaluate( SafetyLevel::FilterAll, &ViewerFeatures::default(), @@ -265,17 +257,15 @@ mod tests { #[test] fn dmca_media_drops_recommendations_only() { let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { + let candidate = candidate() + .with_tweet_features(TweetFeatures { media: MediaFeature { has_dmca_media: true, ..Default::default() }, ..Default::default() - }, - ..Default::default() - }; + }) + .build(); let timeline_home = policies.evaluate( SafetyLevel::TimelineHome, @@ -296,22 +286,16 @@ mod tests { 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 { + let candidate = candidate() + .with_tweet_features(TweetFeatures { nsfw: NsfwFeature { user: true, admin: false, }, ..Default::default() - }, - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + }) + .build(); + let viewer = viewer(VIEWER_ID); let timeline_home = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -332,29 +316,16 @@ mod tests { #[test] fn nsfw_author_interstitials_in_network_but_drops_oon() { - use crate::models::{AuthorFeatures, TweetFeatures}; + use crate::models::AuthorFeatures; let policies = Policies::new(); - let candidate = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { + let candidate = candidate() + .with_media() + .with_author_features(AuthorFeatures { is_nsfw_user: true, ..Default::default() - }, - tweet_features: TweetFeatures { - media: MediaFeature { - has_media: true, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }; - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - allows_sensitive_media: false, - ..Default::default() - }; + }) + .build(); + let viewer = viewer(VIEWER_ID); let in_network = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -379,23 +350,15 @@ mod tests { #[test] fn egregious_nsfw_does_not_drop() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; + use crate::models::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 tweet_candidate = candidate() + .with_label(SafetyLabelType::EGREGIOUS_NSFW) + .build(); let user_candidate = candidate_with_author_user_label(LabelValue::EGREGIOUS_NSFW, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + let viewer = viewer(VIEWER_ID); for candidate in [&tweet_candidate, &user_candidate] { let in_network = policies @@ -415,111 +378,11 @@ mod tests { } } - #[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:?}" - ); - - let oon = policies - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON SpamHighRecall should drop, got {oon:?}" - ); - } - - #[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 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() - }; + let mut c = candidate().with_label(label).build(); c.relationship.viewer_follows_author = follows; c } @@ -528,10 +391,7 @@ mod tests { 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() - }; + let viewer = viewer(VIEWER_ID); for label in [ SafetyLabelType::FOSNR_HATEFUL_CONDUCT, SafetyLabelType::FOSNR_VIOLENT_SPEECH, @@ -556,10 +416,7 @@ mod tests { fn fosnr_never_drops_author() { use crate::models::SafetyLabelType; let policies = Policies::new(); - let author = ViewerFeatures { - viewer: Viewer::LoggedIn(100), - ..Default::default() - }; + let author = author_viewer(); for label in [ SafetyLabelType::FOSNR_HATEFUL_CONDUCT, SafetyLabelType::FOSNR_VIOLENT_SPEECH, @@ -585,14 +442,8 @@ mod tests { 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() - }; + let viewer = viewer(VIEWER_ID); + let author = author_viewer(); for follows in [true, false] { let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, follows); @@ -633,24 +484,19 @@ mod tests { #[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 { + let candidate = candidate() + .with_tweet_features(TweetFeatures { media: MediaFeature { geo_deny_list: vec!["de".to_string()], ..Default::default() }, ..Default::default() - }, - ..Default::default() - }; + }) + .build(); let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), country_code: Some("de".to_string()), - ..Default::default() + ..viewer(VIEWER_ID) }; let in_network = policies @@ -676,20 +522,10 @@ mod tests { #[test] fn nsfw_text_drops_oon_but_allows_in_network() { - use crate::models::{SafetyLabel, SafetyLabelMap, SafetyLabelType}; + use crate::models::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 candidate = candidate().with_label(SafetyLabelType::NSFW_TEXT).build(); + let viewer = viewer(VIEWER_ID); let in_network = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -716,17 +552,7 @@ mod tests { 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() - }; + let mut c = candidate().with_author_user_label(label).build(); c.relationship.viewer_follows_author = follows; c } @@ -736,10 +562,7 @@ mod tests { 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() - }; + let viewer = viewer(VIEWER_ID); let in_network = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -768,10 +591,7 @@ mod tests { let policies = Policies::new(); let candidate = candidate_with_author_user_label(LabelValue::RECOMMENDATIONS_BLACKLIST, false); - let viewer = ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - }; + let viewer = viewer(VIEWER_ID); let in_network = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -799,10 +619,7 @@ mod tests { 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 viewer = viewer(VIEWER_ID); for level in [ SafetyLevel::TimelineHome, @@ -821,10 +638,7 @@ mod tests { 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 viewer = viewer(VIEWER_ID); let in_network = policies .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) @@ -846,67 +660,4 @@ mod tests { ); 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 in_network = policies - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network NsfwNearPerfect should allow, got {in_network:?}" - ); - - 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 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 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:?}" - ); - } } diff --git a/visibility-filtering/rules/socialgraph_rules.rs b/visibility-filtering/rules/socialgraph_rules.rs index 691c13c7..c9b4707e 100644 --- a/visibility-filtering/rules/socialgraph_rules.rs +++ b/visibility-filtering/rules/socialgraph_rules.rs @@ -10,12 +10,10 @@ impl Rule for ViewerBlocksAuthorRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { + if context.viewer_is_logged_out() { return VfAction::Allow; } - if candidate.relationship.viewer_blocks_author { + if context.viewer_blocks_author() { return VfAction::Drop(FilteredReason::AuthorBlockViewer); } VfAction::Allow @@ -30,12 +28,10 @@ impl Rule for MutedRetweetsRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { + if context.viewer_is_logged_out() { return VfAction::Allow; } - if candidate.is_retweet() && candidate.relationship.viewer_mutes_retweets_from_author { + if context.is_retweet() && context.viewer_mutes_retweets_from_author() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -50,12 +46,10 @@ impl Rule for ViewerMutesAuthorRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if viewer.viewer_id().is_none() { + if context.viewer_is_logged_out() { return VfAction::Allow; } - if candidate.relationship.viewer_mutes_author { + if context.viewer_mutes_author() { return VfAction::Drop(FilteredReason::ViewerMutesAuthor); } VfAction::Allow @@ -70,25 +64,23 @@ impl Rule for DropExclusiveTweetContentRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - let Some(exc) = &candidate.exclusive_content else { + if !context.is_exclusive_tweet() { return VfAction::Allow; - }; + } - let Some(viewer_id) = viewer.viewer_id() else { + if context.viewer_is_logged_out() { return VfAction::Drop(FilteredReason::ExclusiveTweet); - }; + } - if viewer_id == exc.conversation_author_id { + if context.viewer_is_conversation_author() { return VfAction::Allow; } - if exc.viewer_super_follows_author { + if context.viewer_super_follows_author() { return VfAction::Allow; } - if !candidate.is_retweet() && context.is_author_viewer() { + if !context.is_retweet() && context.is_author_viewer() { return VfAction::Allow; } @@ -100,59 +92,27 @@ impl Rule for DropExclusiveTweetContentRule { mod tests { use super::*; use crate::models::{ - ExclusiveContentFeatures, HydratedTweetCandidate, TweetFeatures, Viewer, - ViewerAuthorRelationship, ViewerFeatures, + ExclusiveContentFeatures, HydratedTweetCandidate, ViewerAuthorRelationship, }; - - fn viewer(id: u64) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(id), - ..Default::default() - } - } - - fn logged_out_viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedOut, - ..Default::default() - } - } + use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; 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() - } + 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 } fn candidate_with_relationship( relationship: ViewerAuthorRelationship, ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - relationship, - ..Default::default() - } + candidate().with_relationship(relationship).build() } #[test] @@ -162,7 +122,7 @@ mod tests { ..Default::default() }); assert!(matches!( - ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::AuthorBlockViewer) )); } @@ -171,7 +131,7 @@ mod tests { 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)), + ViewerBlocksAuthorRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -195,7 +155,7 @@ mod tests { ..Default::default() }); assert!(matches!( - ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::ViewerMutesAuthor) )); } @@ -204,7 +164,7 @@ mod tests { 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)), + ViewerMutesAuthorRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -223,13 +183,15 @@ mod tests { #[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); + let c = candidate() + .with_relationship(ViewerAuthorRelationship { + viewer_mutes_retweets_from_author: true, + ..Default::default() + }) + .retweet_of(99) + .build(); assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::UnspecifiedReason) )); } @@ -241,18 +203,20 @@ mod tests { ..Default::default() }); assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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); + let c = candidate() + .with_relationship(ViewerAuthorRelationship { + viewer_mutes_retweets_from_author: true, + ..Default::default() + }) + .retweet_of(99) + .build(); assert!(matches!( MutedRetweetsRule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), VfAction::Allow @@ -263,8 +227,8 @@ mod tests { fn non_exclusive_tweet_is_allowed() { let rule = DropExclusiveTweetContentRule; let action = rule.evaluate(&crate::rules::test_context( - &viewer(999), - &non_exclusive_candidate(), + &viewer(VIEWER_ID), + &candidate().build(), )); assert!(matches!(action, VfAction::Allow)); } @@ -287,7 +251,7 @@ mod tests { 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)); + let action = rule.evaluate(&crate::rules::test_context(&author_viewer(), &candidate)); assert!(matches!(action, VfAction::Allow)); } diff --git a/visibility-filtering/rules/tes_rules.rs b/visibility-filtering/rules/tes_rules.rs index 6f90983a..b65e0ed7 100644 --- a/visibility-filtering/rules/tes_rules.rs +++ b/visibility-filtering/rules/tes_rules.rs @@ -1,6 +1,5 @@ use crate::models::VfAction; use crate::rules::{Rule, RuleContext}; -use xai_core_entities::entities::TakedownReason; use xai_visibility_filtering::models::FilteredReason; pub struct DropStaleTweetsRule; @@ -11,17 +10,10 @@ impl Rule for 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; + if context.is_stale_tweet() && !context.is_retweet() { + return VfAction::Drop(FilteredReason::UnspecifiedReason); } - VfAction::Drop(FilteredReason::UnspecifiedReason) + VfAction::Allow } } @@ -33,7 +25,7 @@ impl Rule for DropLegalTakendownPostRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if viewer_in_withheld_country(context, legal_takedown_country) { + if !context.is_author_viewer() && context.legal_takedown_in_viewer_country() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -48,70 +40,22 @@ impl Rule for DropLocalLawsTakendownPostRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if viewer_in_withheld_country(context, local_laws_takedown_country) { + if !context.is_author_viewer() && context.local_laws_takedown_in_viewer_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)) - { + if context.media_restricted_in_viewer_country() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -126,8 +70,7 @@ impl Rule for DropTweetsWithDmcaMediaRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if candidate.has_dmca_media() { + if context.has_dmca_media() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -138,16 +81,21 @@ impl Rule for DropTweetsWithDmcaMediaRule { mod tests { use super::*; use crate::models::{ - CoreFeature, HydratedTweetCandidate, MediaFeature, TakedownFeature, TweetFeatures, Viewer, - ViewerFeatures, + HydratedTweetCandidate, MediaFeature, TakedownFeature, TweetFeatures, ViewerFeatures, }; - use xai_core_entities::entities::{EditControl, EditControlInitial}; + use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; + use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; - fn viewer() -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(999), - ..Default::default() - } + fn takedown_candidate(reasons: Vec) -> HydratedTweetCandidate { + candidate() + .with_tweet_features(TweetFeatures { + takedown: TakedownFeature { + reasons, + ..Default::default() + }, + ..Default::default() + }) + .build() } fn stale_edit_control() -> Option { @@ -160,16 +108,14 @@ mod tests { #[test] fn stale_edit_drops() { let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { + let c = candidate() + .with_tweet_features(TweetFeatures { edit_control: stale_edit_control(), ..Default::default() - }, - ..Default::default() - }; + }) + .build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -177,12 +123,9 @@ mod tests { #[test] fn non_stale_tweet_allows() { let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; + let c = candidate().build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -190,53 +133,37 @@ mod tests { #[test] fn stale_retweet_allows() { let rule = DropStaleTweetsRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { + let c = candidate() + .with_tweet_features(TweetFeatures { edit_control: stale_edit_control(), - core: CoreFeature { - source_tweet_id: Some(99), - ..Default::default() - }, ..Default::default() - }, - ..Default::default() - }; + }) + .retweet_of(99) + .build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } fn viewer_with_country(country: &str) -> ViewerFeatures { ViewerFeatures { - viewer: Viewer::LoggedIn(999), country_code: Some(country.to_string()), - ..Default::default() + ..viewer(VIEWER_ID) } } #[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() + let c = takedown_candidate(vec![ + TakedownReason::LegalRequest { + country_code: "de".to_string(), }, - ..Default::default() - }; + TakedownReason::UnspecifiedReason { + country_code: "fr".to_string(), + }, + ]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Drop(_) @@ -246,19 +173,9 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), VfAction::Allow @@ -268,21 +185,11 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -290,19 +197,9 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "de".to_string(), + }]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Allow @@ -312,24 +209,14 @@ mod tests { #[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() + let c = takedown_candidate(vec![ + TakedownReason::BystanderReport { + country_code: "de".to_string(), }, - ..Default::default() - }; + TakedownReason::BystanderReport { + country_code: "fr".to_string(), + }, + ]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("fr"), &c)), VfAction::Drop(_) @@ -339,19 +226,9 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "de".to_string(), + }]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("us"), &c)), VfAction::Allow @@ -361,19 +238,9 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Allow @@ -383,20 +250,10 @@ mod tests { #[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() - }; + let mut c = takedown_candidate(vec![TakedownReason::LegalRequest { + country_code: "de".to_string(), + }]); + c.author_id = VIEWER_ID; assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Allow @@ -406,20 +263,10 @@ mod tests { #[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() - }; + let mut c = takedown_candidate(vec![TakedownReason::BystanderReport { + country_code: "de".to_string(), + }]); + c.author_id = VIEWER_ID; assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Allow @@ -428,21 +275,11 @@ mod tests { #[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() - }; + let c = takedown_candidate(vec![ + TakedownReason::Dmca, + TakedownReason::HatefulImagery, + TakedownReason::Unknown, + ]); assert!(matches!( DropLegalTakendownPostRule .evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), @@ -456,19 +293,16 @@ mod tests { } fn geo_candidate(allow: &[&str], deny: &[&str]) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { + 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() - }, - ..Default::default() - } + }) + .build() } #[test] @@ -545,7 +379,7 @@ mod tests { let rule = DropTweetsWithGeoRestrictedMediaRule; let c = geo_candidate(&["us"], &[]); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -555,7 +389,7 @@ mod tests { let rule = DropTweetsWithGeoRestrictedMediaRule; let c = geo_candidate(&[], &["xx"]); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -565,7 +399,7 @@ mod tests { let rule = DropTweetsWithGeoRestrictedMediaRule; let c = geo_candidate(&[], &["de"]); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -574,7 +408,7 @@ mod tests { fn geo_restricted_drops_even_for_author() { let rule = DropTweetsWithGeoRestrictedMediaRule; let mut c = geo_candidate(&[], &["de"]); - c.author_id = 999; + c.author_id = VIEWER_ID; assert!(matches!( rule.evaluate(&crate::rules::test_context(&viewer_with_country("de"), &c)), VfAction::Drop(_) @@ -595,19 +429,17 @@ mod tests { #[test] fn dmca_drops() { let rule = DropTweetsWithDmcaMediaRule; - let c = HydratedTweetCandidate { - tweet_id: 1, - tweet_features: TweetFeatures { + let c = candidate() + .with_tweet_features(TweetFeatures { media: MediaFeature { has_dmca_media: true, ..Default::default() }, ..Default::default() - }, - ..Default::default() - }; + }) + .build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } diff --git a/visibility-filtering/rules/tweet_flag_rules.rs b/visibility-filtering/rules/tweet_flag_rules.rs index 0738c240..18dac8f7 100644 --- a/visibility-filtering/rules/tweet_flag_rules.rs +++ b/visibility-filtering/rules/tweet_flag_rules.rs @@ -1,18 +1,18 @@ -use crate::models::{TweetFeatures, VfAction}; +use crate::models::VfAction; use crate::rules::{Rule, RuleContext}; use xai_visibility_filtering::models::FilteredReason; #[derive(Clone)] pub struct TweetFlagDropRule { name: &'static str, - flag: fn(&TweetFeatures) -> bool, + flag: fn(&RuleContext<'_>) -> bool, reason: FilteredReason, } impl TweetFlagDropRule { pub const fn new( name: &'static str, - flag: fn(&TweetFeatures) -> bool, + flag: fn(&RuleContext<'_>) -> bool, reason: FilteredReason, ) -> Self { Self { name, flag, reason } @@ -25,7 +25,7 @@ impl Rule for TweetFlagDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if (self.flag)(&context.candidate().tweet_features) { + if (self.flag)(context) { return VfAction::Drop(self.reason.clone()); } VfAction::Allow @@ -34,46 +34,35 @@ impl Rule for TweetFlagDropRule { pub const TWEET_NSFW_USER_DROP: TweetFlagDropRule = TweetFlagDropRule::new( "TweetNsfwUserDropRule", - |t| t.nsfw.user, + |context| context.has_tweet_nsfw_user_flag(), FilteredReason::ContainNsfwMedia, ); pub const TWEET_NSFW_ADMIN_DROP: TweetFlagDropRule = TweetFlagDropRule::new( "TweetNsfwAdminDropRule", - |t| t.nsfw.admin, + |context| context.has_tweet_nsfw_admin_flag(), FilteredReason::ContainNsfwMedia, ); #[cfg(test)] mod tests { use super::*; - use crate::models::{ - HydratedTweetCandidate, NsfwFeature, TweetFeatures, Viewer, ViewerFeatures, - }; + use crate::models::{HydratedTweetCandidate, NsfwFeature, TweetFeatures}; + use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; fn candidate_with_nsfw_flags(user: bool, admin: bool) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - tweet_features: TweetFeatures { + candidate() + .with_tweet_features(TweetFeatures { nsfw: NsfwFeature { user, admin }, ..Default::default() - }, - ..Default::default() - } - } - - fn viewer(id: u64) -> ViewerFeatures { - ViewerFeatures { - viewer: Viewer::LoggedIn(id), - ..Default::default() - } + }) + .build() } #[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)), + TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } @@ -82,7 +71,7 @@ mod tests { 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)), + TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -91,7 +80,7 @@ mod tests { 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)), + TWEET_NSFW_USER_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } @@ -100,7 +89,7 @@ mod tests { 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)), + TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } @@ -109,7 +98,7 @@ mod tests { 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)), + TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -118,7 +107,7 @@ mod tests { 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)), + TWEET_NSFW_ADMIN_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } diff --git a/visibility-filtering/rules/tweet_label_drops.rs b/visibility-filtering/rules/tweet_label_drops.rs index d5c9b2ab..a915760f 100644 --- a/visibility-filtering/rules/tweet_label_drops.rs +++ b/visibility-filtering/rules/tweet_label_drops.rs @@ -158,71 +158,49 @@ pub const FOSNR_ABUSE_INSULTS_OON_DROP: SafetyLabelDropRule = SafetyLabelDropRul #[cfg(test)] mod tests { use super::*; - use crate::models::{ - HydratedTweetCandidate, SafetyLabel, SafetyLabelMap, Viewer, ViewerFeatures, + use crate::rules::fixtures::{ + author_viewer, candidate, sensitive_opt_in_viewer, viewer, VIEWER_ID, }; - 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); + let c = candidate().with_label(SafetyLabelType::PDNA).build(); assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::SafetyResult(_)) )); - let c = candidate_with_label(SafetyLabelType::BOUNCE); + let c = candidate().with_label(SafetyLabelType::BOUNCE).build(); assert!(matches!( - BOUNCE_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + BOUNCE_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::TweetIsBounced) )); - let c = candidate_with_label(SafetyLabelType::SPAM); + let c = candidate().with_label(SafetyLabelType::SPAM).build(); assert!(matches!( - SPAM_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + SPAM_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::PossiblyUndesirable) )); - let c = candidate_with_label(SafetyLabelType::NSFW_HIGH_RECALL); + let c = candidate() + .with_label(SafetyLabelType::NSFW_HIGH_RECALL) + .build(); assert!(matches!( - NSFW_HIGH_RECALL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + NSFW_HIGH_RECALL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } #[test] fn author_exempt_rules_allow_self_view() { - let c = candidate_with_label(SafetyLabelType::PDNA); + let c = candidate().with_label(SafetyLabelType::PDNA).build(); assert!(matches!( PDNA_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); - let c = candidate_with_label(SafetyLabelType::DO_NOT_AMPLIFY); + let c = candidate() + .with_label(SafetyLabelType::DO_NOT_AMPLIFY) + .build(); assert!(matches!( DO_NOT_AMPLIFY_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow @@ -231,7 +209,9 @@ mod tests { #[test] fn all_viewer_rules_drop_even_for_author() { - let c = candidate_with_label(SafetyLabelType::FOR_EMERGENCY_USE_ONLY); + let c = candidate() + .with_label(SafetyLabelType::FOR_EMERGENCY_USE_ONLY) + .build(); assert!(matches!( FOR_EMERGENCY_USE_ONLY_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Drop(FilteredReason::UnspecifiedReason) @@ -240,43 +220,25 @@ mod tests { #[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); + let c = candidate() + .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) + .build(); assert!(matches!( - NSFW_HIGH_PRECISION_DROP.evaluate(&crate::rules::test_context(&opted_in, &c)), + NSFW_HIGH_PRECISION_DROP + .evaluate(&crate::rules::test_context(&sensitive_opt_in_viewer(), &c)), VfAction::Drop(FilteredReason::ContainNsfwMedia) )); } #[test] fn no_label_allows() { - let c = HydratedTweetCandidate { - tweet_id: 1, - ..Default::default() - }; + let c = candidate().build(); assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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 [ @@ -285,14 +247,14 @@ mod tests { &FOSNR_ABUSE_DROP, &FOSNR_CIVIC_INTEGRITY_DROP, ] { - let c = candidate_with_label(rule.label); + let c = candidate().with_label(rule.label).build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::PossiblyUndesirable) )); - let c = candidate_with_label_followed(rule.label); + let c = candidate().with_label(rule.label).followed().build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&follower_viewer(), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::PossiblyUndesirable) )); } @@ -306,7 +268,7 @@ mod tests { &FOSNR_ABUSE_DROP, &FOSNR_CIVIC_INTEGRITY_DROP, ] { - let c = candidate_with_label(rule.label); + let c = candidate().with_label(rule.label).build(); assert!(matches!( rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow @@ -316,39 +278,46 @@ mod tests { #[test] fn malicious_url_drops_non_author_and_exempts_author() { - let c = candidate_with_label(SafetyLabelType::MALICIOUS_URL); + let c = candidate() + .with_label(SafetyLabelType::MALICIOUS_URL) + .build(); assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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() - }; + let c = candidate().build(); assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } #[test] fn fosnr_abuse_insults_oon_drops_all_non_authors() { - let c = candidate_with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS); + let c = candidate() + .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) + .build(); assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP.evaluate(&crate::rules::test_context(&viewer(), &c)), + FOSNR_ABUSE_INSULTS_OON_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::PossiblyUndesirable) )); - let c = candidate_with_label_followed(SafetyLabelType::FOSNR_ABUSE_INSULTS); + let c = candidate() + .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) + .followed() + .build(); assert!(matches!( FOSNR_ABUSE_INSULTS_OON_DROP - .evaluate(&crate::rules::test_context(&follower_viewer(), &c)), + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::PossiblyUndesirable) )); - let c = candidate_with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS); + let c = candidate() + .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) + .build(); assert!(matches!( FOSNR_ABUSE_INSULTS_OON_DROP .evaluate(&crate::rules::test_context(&author_viewer(), &c)), diff --git a/visibility-filtering/rules/user_label_drops.rs b/visibility-filtering/rules/user_label_drops.rs index 54b87a8e..09b6fdd1 100644 --- a/visibility-filtering/rules/user_label_drops.rs +++ b/visibility-filtering/rules/user_label_drops.rs @@ -36,11 +36,11 @@ impl Rule for UserSafetyLabelDropRule { if context.is_author_viewer() { return VfAction::Allow; } - if !context.candidate().author_has_user_label(self.label) { + if !context.author_has_user_label(self.label) { return VfAction::Allow; } if self.require_non_follower - && !context.viewer().viewer_is_logged_out() + && !context.viewer_is_logged_out() && context.viewer_follows_author() { return VfAction::Allow; @@ -121,54 +121,32 @@ pub const DO_NOT_AMPLIFY_NON_FOLLOWER_USER_DROP: UserSafetyLabelDropRule = #[cfg(test)] mod tests { use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, UserLabelSet, Viewer, ViewerFeatures, - }; - use std::collections::HashSet; + use crate::models::HydratedTweetCandidate; + use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; 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() - } + candidate().with_author_user_label(label).build() } #[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)), + NSFW_HIGH_RECALL_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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)), + COMPROMISED_USER_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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)), + SPAM_HIGH_RECALL_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::UnspecifiedReason) )); } @@ -197,14 +175,10 @@ mod tests { #[test] fn allows_when_label_absent() { - let c = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; + let c = candidate().build(); assert!(matches!( IMPERSONATION_HIGH_PRECISION_USER_DROP - .evaluate(&crate::rules::test_context(&viewer(), &c)), + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -213,7 +187,8 @@ mod tests { 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)), + NSFW_HIGH_PRECISION_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -222,13 +197,15 @@ mod tests { 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)), + NSFW_AVATAR_IMAGE_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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)), + NSFW_BANNER_IMAGE_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(FilteredReason::UnspecifiedReason) )); } @@ -237,16 +214,14 @@ mod tests { 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)), + ABUSIVE_HIGH_RECALL_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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)), + ABUSIVE_HIGH_RECALL_USER_DROP + .evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), VfAction::Drop(FilteredReason::UnspecifiedReason) )); } @@ -256,7 +231,8 @@ mod tests { 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)), + ABUSIVE_HIGH_RECALL_USER_DROP + .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } diff --git a/visibility-filtering/rules/user_rules.rs b/visibility-filtering/rules/user_rules.rs index 893ddd03..b00f6d9b 100644 --- a/visibility-filtering/rules/user_rules.rs +++ b/visibility-filtering/rules/user_rules.rs @@ -1,18 +1,18 @@ -use crate::models::{AuthorFeatures, VfAction}; +use crate::models::VfAction; use crate::rules::{Rule, RuleContext}; use xai_visibility_filtering::models::FilteredReason; #[derive(Clone)] pub struct AuthorFlagDropRule { name: &'static str, - flag: fn(&AuthorFeatures) -> bool, + flag: fn(&RuleContext<'_>) -> bool, reason: FilteredReason, } impl AuthorFlagDropRule { pub const fn new( name: &'static str, - flag: fn(&AuthorFeatures) -> bool, + flag: fn(&RuleContext<'_>) -> bool, reason: FilteredReason, ) -> Self { Self { name, flag, reason } @@ -25,8 +25,7 @@ impl Rule for AuthorFlagDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let candidate = context.candidate(); - if (self.flag)(&candidate.author_features) && !context.is_author_viewer() { + if (self.flag)(context) && !context.is_author_viewer() { return VfAction::Drop(self.reason.clone()); } VfAction::Allow @@ -35,32 +34,32 @@ impl Rule for AuthorFlagDropRule { pub const SUSPENDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "SuspendedAuthorRule", - |a| a.is_suspended, + |context| context.author_is_suspended(), FilteredReason::AuthorIsSuspended, ); pub const DEACTIVATED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DeactivatedAuthorRule", - |a| a.is_deactivated, + |context| context.author_is_deactivated(), FilteredReason::AuthorIsDeactivated, ); pub const ERASED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "ErasedAuthorRule", - |a| a.is_erased, + |context| context.author_is_erased(), FilteredReason::AuthorAccountIsInactive, ); pub const OFFBOARDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "OffboardedAuthorRule", - |a| a.is_offboarded, + |context| context.author_is_offboarded(), FilteredReason::AuthorAccountIsInactive, ); pub const NSFW_USER_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DropNsfwUserAuthorRule", - |a| a.is_nsfw_user, + |context| context.author_is_nsfw_user(), FilteredReason::ContainNsfwMedia, ); pub const NSFW_ADMIN_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DropNsfwAdminAuthorRule", - |a| a.is_nsfw_admin, + |context| context.author_is_nsfw_admin(), FilteredReason::ContainNsfwMedia, ); @@ -72,11 +71,9 @@ impl Rule for ProtectedAuthorDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - let viewer = context.viewer(); - let candidate = context.candidate(); - if candidate.author_features.is_protected + if context.author_is_protected() && !context.is_author_viewer() - && (viewer.viewer_is_logged_out() || !candidate.viewer_follows_author()) + && (context.viewer_is_logged_out() || !context.viewer_follows_author()) { return VfAction::Drop(FilteredReason::AuthorIsProtected); } @@ -87,33 +84,22 @@ impl Rule for ProtectedAuthorDropRule { #[cfg(test)] mod tests { use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, Viewer, ViewerAuthorRelationship, ViewerFeatures, - }; + use crate::models::{AuthorFeatures, HydratedTweetCandidate}; + use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; fn candidate_with_author( suspended: bool, deactivated: bool, protected: bool, ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { + candidate() + .with_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() - } + }) + .build() } #[test] @@ -121,7 +107,7 @@ mod tests { let rule = SUSPENDED_AUTHOR_DROP; let c = candidate_with_author(true, false, false); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -131,7 +117,7 @@ mod tests { let rule = SUSPENDED_AUTHOR_DROP; let c = candidate_with_author(true, false, false); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), + rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } @@ -141,7 +127,7 @@ mod tests { let rule = DEACTIVATED_AUTHOR_DROP; let c = candidate_with_author(false, true, false); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -151,7 +137,7 @@ mod tests { let rule = ProtectedAuthorDropRule; let c = candidate_with_author(false, false, true); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -160,12 +146,9 @@ mod tests { 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() - }; + c.relationship.viewer_follows_author = true; assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -174,12 +157,8 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &c)), VfAction::Drop(FilteredReason::AuthorIsProtected) )); } @@ -189,7 +168,7 @@ mod tests { let rule = ProtectedAuthorDropRule; let c = candidate_with_author(false, false, true); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(100), &c)), + rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } @@ -216,23 +195,14 @@ mod tests { ]; 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() - }; + let flagged = candidate().with_author_features(author_features).build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &flagged)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &flagged)), VfAction::Drop(FilteredReason::AuthorAccountIsInactive) )); - let unflagged = HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - ..Default::default() - }; + let unflagged = candidate().build(); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &unflagged)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &unflagged)), VfAction::Allow )); } @@ -243,7 +213,7 @@ mod tests { let rule = SUSPENDED_AUTHOR_DROP; let c = candidate_with_author(false, false, false); assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(999), &c)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } @@ -252,16 +222,13 @@ mod tests { is_nsfw_user: bool, is_nsfw_admin: bool, ) -> HydratedTweetCandidate { - HydratedTweetCandidate { - tweet_id: 1, - author_id: 100, - author_features: AuthorFeatures { + candidate() + .with_author_features(AuthorFeatures { is_nsfw_user, is_nsfw_admin, ..Default::default() - }, - ..Default::default() - } + }) + .build() } #[test] @@ -269,7 +236,7 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -279,7 +246,7 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } @@ -289,7 +256,7 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Drop(_) )); } @@ -299,7 +266,7 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), VfAction::Allow )); } @@ -309,7 +276,7 @@ mod tests { 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)), + rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), VfAction::Allow )); } From 6384ca7d2c8570fbc645c20c3291730739ac00ce Mon Sep 17 00:00:00 2001 From: CI agent Date: Tue, 1 Sep 2026 00:11:04 +0000 Subject: [PATCH 12/18] Open-source X Recommendation Algorithm --- .../service-lib/src/config.rs | 125 ++------ .../service-lib/src/lib.rs | 152 ++++----- grox/config/config.py | 1 + grox/core/data_loaders/media_loader.py | 20 ++ grox/core/lm/convo.py | 29 +- grox/flows/ptos/constants.py | 4 + grox/flows/ptos/generators.py | 12 + grox/flows/ptos/mode.py | 2 +- .../ptos/plan_safety_ptos_special_video.py | 55 ++++ grox/flows/ptos/task_rate_limit.py | 4 +- ...ety_ptos_adult_content_cross_validation.py | 27 +- .../task_safety_ptos_safemodel_sex_nudity.py | 4 +- .../task_safety_ptos_special_video_filter.py | 40 +++ grox/flows/ptos/task_special_video_screen.py | 49 +++ grox/libs/video_tools/image.py | 125 ++++++++ grox/libs/video_tools/video_frames.py | 53 +++- grox/main.py | 73 +++++ .../vf_candidate_hydrator.rs | 12 +- .../vf_following_candidate_hydrator.rs | 3 + .../filters/brazil_2026_election_filter.rs | 296 +++++++++++++----- home-mixer/models/candidate.rs | 7 +- home-mixer/params/param.rs | 8 +- .../phoenix_request_cache_side_effect.rs | 5 +- .../side_effects/scored_stats_side_effect.rs | 76 ++++- phoenix/crates/common/xai-recsys/Cargo.toml | 1 + .../common/xai-recsys/src/model_config.rs | 98 +++++- phoenix/crates/common/xai-recsys/src/util.rs | 7 +- .../serving/xai-recsys-engine/src/python.rs | 2 +- .../serving/xai-recsys-mm-server/Cargo.toml | 1 + .../src/mm_embedding_client.rs | 183 +++++++---- .../xai-recsys-mm-server/src/snapshot.rs | 50 ++- .../xai-recsys-proto/proto/recsys.proto | 82 ++--- .../common/xai-proto/proto/recsys.proto | 82 ++--- .../python/common/xai-proto/pyproject.toml | 4 +- .../xai_checkpointing/dek.py | 49 +++ .../xai_checkpointing/encrypted_kvstore.py | 49 +++ phoenix/xrex/configs/data_feeds.py | 2 + phoenix/xrex/configs/xrecsys_two_tower.py | 8 +- phoenix/xrex/driver/config_factory.py | 2 +- phoenix/xrex/driver/{driver.py => core.py} | 0 .../xrex/driver/{driver_local.py => local.py} | 2 +- phoenix/xrex/utils/checkpointing.py | 5 +- .../EntityToSimClustersEmbeddingsJob.scala | 19 +- thunder/args.rs | 3 + thunder/kafka_utils.rs | 105 ++++++- .../tweet_safety_label.rs | 2 +- visibility-filtering-client/vf_client.rs | 161 ++++++++-- visibility-filtering/config.rs | 7 + visibility-filtering/filter.rs | 18 +- visibility-filtering/filter_tweets.rs | 47 ++- visibility-filtering/lib.rs | 1 + visibility-filtering/main.rs | 1 + visibility-filtering/params.rs | 216 +++++++++++++ visibility-filtering/reference_compare.rs | 39 ++- visibility-filtering/rules/context.rs | 273 ++++++++++------ visibility-filtering/rules/metrics.rs | 12 + visibility-filtering/rules/nsfw_age_gating.rs | 76 +++-- .../rules/nsfw_interstitial.rs | 14 +- visibility-filtering/rules/nullcast_rule.rs | 5 +- visibility-filtering/rules/registry.rs | 64 +++- .../rules/socialgraph_rules.rs | 22 +- visibility-filtering/rules/tes_rules.rs | 10 +- .../rules/tweet_flag_rules.rs | 4 +- .../rules/tweet_label_drops.rs | 4 +- .../rules/user_label_drops.rs | 8 +- visibility-filtering/rules/user_rules.rs | 20 +- visibility-filtering/server.rs | 9 +- visibility-filtering/server_deps.rs | 11 +- visibility-filtering/twemcache/host_pool.rs | 16 +- visibility-filtering/twemcache/metrics.rs | 6 +- 70 files changed, 2247 insertions(+), 735 deletions(-) create mode 100644 grox/flows/ptos/plan_safety_ptos_special_video.py create mode 100644 grox/flows/ptos/task_safety_ptos_special_video_filter.py create mode 100644 grox/flows/ptos/task_special_video_screen.py create mode 100644 grox/main.py create mode 100644 phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py create mode 100644 phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py rename phoenix/xrex/driver/{driver.py => core.py} (100%) rename phoenix/xrex/driver/{driver_local.py => local.py} (98%) create mode 100644 visibility-filtering/params.rs 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/lib.rs b/abuse-enforcement-service/service-lib/src/lib.rs index e3060684..991a4265 100644 --- a/abuse-enforcement-service/service-lib/src/lib.rs +++ b/abuse-enforcement-service/service-lib/src/lib.rs @@ -42,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; @@ -663,25 +662,9 @@ async fn run_enforcement( 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)"); @@ -691,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 { @@ -1441,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}; @@ -1479,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 = @@ -1539,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, @@ -1567,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 @@ -1586,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:#}" )); } } @@ -1790,24 +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_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); + .with_fetch_timeout_ms(10000) + .build(); - BatchConsumerConfig::new(kafka, SERVICE_NAME) - .with_max_messages_per_poll(max_messages_per_poll) + 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 { @@ -1877,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) @@ -1948,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; diff --git a/grox/config/config.py b/grox/config/config.py index 6f9575ea..3fd6a690 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -110,6 +110,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 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/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/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_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/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 index afd886fe..648167d9 100644 --- a/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py +++ b/grox/flows/ptos/task_safety_ptos_adult_content_cross_validation.py @@ -5,7 +5,7 @@ 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 SAFETY_PTOS_DELUXE +from grox.flows.ptos.constants import DELUXE_TIER_TASK_TYPES, SAFETY_PTOS_SPECIAL_VIDEO from grox.flows.ptos.state import ( SafetyPolicy, SafetyPolicyCategory, @@ -20,6 +20,7 @@ _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): @@ -55,7 +56,7 @@ async def _run(cls, ctx: TaskContext, post: Post) -> None: if not state.safemodel_sex_nudity.scored: return - 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" safemodel_positive = state.safemodel_sex_nudity.positive violations = ( @@ -104,9 +105,15 @@ async def _cross_validate(cls, ctx: TaskContext, post: Post) -> None: judged = await cls._judge.judge(post) except Exception as e: Metrics.counter(metric).add(1, attributes={"outcome": "error"}) - logger.warning( - f"Post {post.id}: grok 4.5 cross validation failed, keeping original decision: {e}" - ) + 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 @@ -117,6 +124,10 @@ async def _cross_validate(cls, ctx: TaskContext, post: Post) -> None: 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 = [ @@ -125,15 +136,13 @@ async def _cross_validate(cls, ctx: TaskContext, post: Post) -> None: if not adult_violations: adult_violations = [ SafetyPtosViolatedPolicy( - category=SafetyPolicyCategory.AdultContent, - reason=_CROSS_VALIDATION_REASON, + category=SafetyPolicyCategory.AdultContent, reason=reason ) ] violations.append(adult_violations[0]) for violation in adult_violations: violation.safetyPolicy = SafetyPolicy( - policyType=SafetyPolicyType.AdultContentSexualSoft, - reason=_CROSS_VALIDATION_REASON, + policyType=SafetyPolicyType.AdultContentSexualSoft, reason=reason ) state.annotations.violatedPolicies = violations state.safemodel_sex_nudity.positive = False 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 7a4417b1..aafea867 100644 --- a/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py +++ b/grox/flows/ptos/task_safety_ptos_safemodel_sex_nudity.py @@ -13,7 +13,7 @@ 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__) @@ -59,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): 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/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/candidate_hydrators/vf_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs index 1d935657..264a7a8e 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(); } @@ -104,8 +104,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 { diff --git a/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs index c32e69dc..19273fc0 100644 --- a/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs +++ b/home-mixer/candidate_hydrators/vf_following_candidate_hydrator.rs @@ -63,6 +63,9 @@ impl Hydrator for VFFollowingCandidateHydrator 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()); diff --git a/home-mixer/filters/brazil_2026_election_filter.rs b/home-mixer/filters/brazil_2026_election_filter.rs index 78d92e4b..1d0886ee 100644 --- a/home-mixer/filters/brazil_2026_election_filter.rs +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -19,6 +19,7 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // 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 @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. @@ -26,7 +27,6 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // 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. -// @ADALBERTO_1111 no live account found. // @ADRIANASOUSAPIAUI no live account found. // @AGOLDBACH no live account found. // @AHELIXO no live account found. @@ -36,6 +36,7 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @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. @@ -46,6 +47,7 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @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. @@ -64,14 +66,14 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @MIRCOCORONETTI no live account found. // @NADIAGERHARD no live account found. // @NETOFEITOSA6891 no live account found. -// @OMARAZIZSENADOR no live account found. // @PATRICIACRIZANTO2 no live account found. // @PAULOMOURAOTO no live account found. // @PEDRONASSIF_RJ no live account found. // @PEDROPONCIOBE no live account found. // @POLICIALPAULOBASTOS no live account found. // @PRADOCORONEL no live account found. -// @SENATORCIDGOMES 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. @@ -88,6 +90,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 14492205, // @pedro_lupion 15022409, + // @erasintetica + 15585094, // @soninhafrancine 15768105, // @tatyanavaleria @@ -196,8 +200,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 34430921, // @RafaCupertino 34618485, - // @carlosviana - 34630924, // @BetoRicha 34665220, // @Donato_PT @@ -242,6 +244,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 37700244, // @dep_geraldo 37711911, + // @SouCrisCaiado + 37745764, // @ruialves10_ 37949658, // @mauricioscalco @@ -286,6 +290,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 42455446, // @kruke1 42487937, + // @Domingos_Neto + 42626592, // @ArlenSantiago 42630237, // @Francischini_ @@ -296,6 +302,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 43189774, // @axelgrael 43326346, + // @owagnerpro + 43364776, // @edurodrigues_25 43856097, // @profleomatos @@ -310,6 +318,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 44455876, // @rafaelcampelo 44460118, + // @rafamacris + 44690902, // @perpetua_acre 44693900, // @pedrosuplicy30 @@ -342,8 +352,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 47461491, // @augustocury 47529845, - // @naderaliumar - 47655960, // @Altineu 47991805, // @fernandojordao @@ -434,6 +442,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 52842057, // @RodrigoGuedesam 52954632, + // @joliveirambl + 52971495, // @RollembergPSB 53050115, // @jorginhomello @@ -526,14 +536,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 57151676, // @HarifeViegas 57163107, - // @GustinhoRibeiro - 57208702, // @Rafaelpicciani 57529926, // @CanzianiAlex 57641073, // @depHugoLeal 57771926, + // @gabriellimambl + 58029195, + // @KarenGregoris + 58044789, // @marinapassadore 58247896, // @AlmeidaMarcus @@ -580,6 +592,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 60731692, // @Isquierdorio 60805457, + // @luisbenambl + 60983130, // @CovattiFilho 60994156, // @franzepiaui @@ -752,12 +766,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 70956209, // @rdlorenzoni 71056246, + // @sen_wellington + 71065638, // @charlesribeiro_ 71098452, // @lindberghfarias 71310152, - // @lincolndrumond - 71541588, // @Daniel_PCdoB 71545154, // @FaissalCalil @@ -798,6 +812,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 74361905, // @LeurLomantoJr 74538721, + // @RogerPerestelo + 74583511, // @Casagrande_ES 74722174, // @luizcoutopt @@ -862,6 +878,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 78714361, // @profdorinha 79174387, + // @gilbertoabramo + 80095058, // @michelschlemper 80123403, // @capitaotadeu @@ -936,8 +954,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 84608790, // @filhomarcio 84639445, - // @giselecasarin - 84943825, // @lubloureiro 85150664, // @RobertoPSOL @@ -946,6 +962,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 85327394, // @emanuelcacho 85461555, + // @danielter + 85585608, // @maxlemos 85613796, // @PCBpartidao @@ -984,6 +1002,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 90898890, // @wilsonlimaAM 91109801, + // @RaffieDellon + 91840828, // @f_trad 92509126, // @felipeaugusto01 @@ -1046,6 +1066,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 100521721, // @julioarcoverde 101016613, + // @fernando_it + 101892949, // @PabloValenteDF 102257530, // @EduardoGomesTO @@ -1084,8 +1106,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 108988113, // @rosenvergreis 109006854, - // @CharlesDrumond - 109041140, // @DiogoForjaz 109147422, // @deplucasdelima @@ -1196,8 +1216,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 128550627, // @GodriJunior 128906180, - // @Danusalopes - 129028918, // @apjunqueira 129055364, // @ninamarinabraga @@ -1232,6 +1250,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 136381913, // @paulomansur_ 136710714, + // @andreia_zito + 137369968, // @BiradoPindare 137548563, // @tadeuveneri @@ -1312,6 +1332,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 158067709, // @Marivaldo4P 159643822, + // @CoronelMedina14 + 160541377, // @sandroalexpr 160554168, // @MarcosRogerio @@ -1388,10 +1410,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 168657354, // @pepecollaco 169158984, - // @neidermoreira - 169515284, // @matheusmanholer 170176086, + // @pedrofrancez + 170188044, // @walterlfcaval 170638771, // @wilsonsousajr @@ -1448,6 +1470,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 185482131, // @rodfvale 186331377, + // @acarlosmendes + 190308328, // @pablomarcal 191223319, // @lorran_rebeldia @@ -1456,8 +1480,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 192658045, // @pretagilsa 193045733, - // @analacerdamg - 193138273, // @MarciaTaschetti 198338329, // @JacksonAndre7 @@ -1504,8 +1526,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 213382232, // @rafaelprimo 215727649, - // @DanielBarbosaAL - 217322775, // @Lucinildo 217347443, // @celmarcosantos @@ -1562,6 +1582,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 245320150, // @RogerioCorreia_ 245392082, + // @gelizabetesp + 247745869, // @DeAssisDiniz 247906787, // @iginomarcos13 @@ -1640,8 +1662,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 285614672, // @maicolmed 286427700, - // @RenataMelorj - 286972716, // @gabrielaorttiz 287219048, // @marcelomaranata @@ -1694,6 +1714,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 307557586, // @gleicejanems 308780059, + // @guto_schiavetto + 309004600, // @SamiraDaud 310033093, // @GildeteAlves @@ -1778,8 +1800,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 370942708, // @RONIESILVA15 373465468, - // @AdailtonAdvog - 379850268, // @coroneljunior 381045128, // @ieda_chaves @@ -1814,8 +1834,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 423874716, // @vitordeangelo 424455116, - // @luisfe_valdivia - 427885559, // @BalbinottiFilho 428247512, // @RealNabor @@ -1848,6 +1866,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 465613391, // @bellagoncalvs 469918968, + // @crisbrasilreal + 472033751, // @andreawerner_ 475996406, // @leopratesba @@ -1862,8 +1882,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 487622592, // @LPescinelli 492334281, - // @brunosouzasc - 494268633, // @brauliolaranovo 505339844, // @glauberbastos_ @@ -1876,8 +1894,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 554803332, // @alcyvania 556852639, + // @limmajr168 + 565889975, // @ThammyReal 578495086, + // @mariohildebrand + 580176611, // @OthelinoNeto 583377940, // @HelioWirbiski @@ -1922,6 +1944,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 799260530, // @JenirNeves 813802178, + // @PedroDeyrot + 857054846, // @carlaopelobem 893975196, // @Isoldadantaspt @@ -1972,8 +1996,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1356677952, // @D_GoretePereira 1362596354, - // @VAGNERVISOLI - 1420675674, // @RobertoRocha_MA 1436541721, // @mickasevalho @@ -2150,6 +2172,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2838953716, // @daianasantospoa 2858823694, + // @yagorVieira + 2868228117, // @_akalicia_ 2879108776, // @paulolemosap @@ -2164,6 +2188,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2925491427, // @KuhlmannJean 2927136550, + // @NoletoMBL + 2938978041, // @DepJuscelino 2970617333, // @neioluciofp @@ -2172,6 +2198,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2977624732, // @guipaoficial 2979670457, + // @depjosenelto + 2995610423, // @isaakalmeida93 2997861520, // @carlosveraspt @@ -2212,6 +2240,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3130887358, // @DrLeonardomt 3131609429, + // @senadorasoraya + 3167874665, // @meire_cruvinel 3205786257, // @Marcio_Honaiser @@ -2228,6 +2258,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3342622547, // @PbnConcursos 3357932231, + // @brasleiroluc + 3366080079, // @moisesbrazpt 3373574517, // @kleybe_morais @@ -2236,12 +2268,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3662771592, // @OPaiakan 3674682197, - // @diegolopesadv - 3734258421, // @delegadanadine 3744465381, - // @heliomissao - 3853530796, // @ninasouzarn 3904595243, // @FelipeMichelRJ @@ -2458,12 +2486,22 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 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 @@ -2494,8 +2532,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 947255546062745600, // @jusmarioficial 948577496177561600, - // @AnaFialho14 - 952537771717033984, // @DpRicardoArruda 953055428124045313, // @Carlos_cabral81 @@ -2508,6 +2544,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 959184931552464896, // @Fbgg40 963928182368980993, + // @RafaelLustoza_ + 965096425313972229, // @XandePessoaPE 966833467391725569, // @Pastorellux @@ -2548,8 +2586,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 984221723544444928, // @eng_angelo44 984522534623301632, - // @AllanAguiar14 - 984581958943629313, // @israelsantosap 985962825909723141, // @Dep_GilPereira @@ -2580,6 +2616,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 993288307625943040, // @andreiadejesuus 994268335486504963, + // @Alessandrojesu_ + 996849375635820544, // @prdinhosouza 999747749364076545, // @thiagoavilabr @@ -2700,14 +2738,20 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1053123917785808901, // @RafaelDemarchi5 1053334763858214912, + // @deputadatalita + 1055640991200411648, // @cleitinhotmj 1057231251743170562, // @delegadasheila 1058010509256126464, + // @mauricio_lindol + 1058341993372360705, // @veronicalima_ve 1059554967600685058, // @pinheirinhomg 1060134845043666945, + // @giordanmes + 1062407588724252673, // @pluviapt 1062505159824150530, // @capitaocarpe @@ -2806,6 +2850,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1095990374550695938, // @DiogoTalento 1096490764677320704, + // @ReMesquita1977 + 1097042489813409792, // @DFDanielFreitas 1097498693719199744, // @dilvandafaroPT @@ -2986,6 +3032,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1177185451536465920, // @marcimeirelles 1177246955942154241, + // @GalbaSheyla + 1179012777366736899, // @delegadopalumbo 1179437585275465729, // @rickazzevedo @@ -3022,8 +3070,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1208038091962888194, // @LeoSuricate 1208544960032727040, - // @franciscodiasup - 1209429703909691392, // @ManuVieiraSC 1210296676520513536, // @CruzOrleans @@ -3036,8 +3082,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1213876635331465216, // @OficialNenemAl 1216746154626625536, - // @BrenoFonsecaMG - 1216897633769443329, // @JairSoutoAM 1218904655591243777, // @JohnRobertPA @@ -3060,6 +3104,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1224592244465905666, // @KlesleyGarcia 1224615230195609601, + // @PortaldoJose + 1224767212101283840, // @joaobmaresguia 1225055364124762112, // @profterezinhaPT @@ -3106,6 +3152,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1238835947661377536, // @eusouamom 1238868377675980803, + // @eduardomenegol_ + 1240359378081001474, // @juliana_macieel 1240694913420926976, // @fabriciochaves_ @@ -3124,6 +3172,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1242093483151839238, // @claudiaguerramg 1242834650155941888, + // @gabrielcarvalce + 1243211951863463941, // @samaramartinsup 1244783574676627460, // @docporto @@ -3260,12 +3310,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1285295088525082631, // @Sonaira_sp 1285299871340167173, + // @euLiviaNoronha + 1285729712795525125, // @ladisouzams 1286438113896796170, // @ericodonovo 1287396715331452932, // @sorriso_elisa 1287490158510723072, + // @prjuniortrovao + 1288544254856507392, // @majorvitorsa 1290713462285500417, // @MatheusLaiola @@ -3382,8 +3436,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1352236209800671232, // @dimasfabianomg 1355216660068761606, - // @SeccoHelio - 1355364623181094914, // @AnaPimentelmg 1356434859934298115, // @faustinorn01 @@ -3422,8 +3474,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1379025327314329608, // @DrLeviMelo1 1379424037068234752, - // @fabiosilveirarn - 1379583456812924935, // @marleipr 1380366998488645636, // @depprofcleiton @@ -3450,8 +3500,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1396055754793230341, // @jusoaresft 1400452986178981891, - // @oemersonmatos - 1400821681619353602, // @eduacostario 1400996132982038532, // @antidio_lunelli @@ -3512,6 +3560,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1439763846772703232, // @edilenxavier 1441394605338021888, + // @ivanirdosantos + 1442471110927290379, // @mariiluse 1445155253880594432, // @renatmirandarj @@ -3626,6 +3676,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1497201650196430861, // @profangelapsol 1497605679229644803, + // @fsdecastilho + 1499529892039434251, // @EnfBrunoFarias 1500910030593445888, // @annasebbaj @@ -3642,6 +3694,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1506313161921777676, // @rhdeverdade 1508246534370082821, + // @bm_oliveira14 + 1508783941037268993, // @Thais_ProfeChef 1509245802815922177, // @MaalouliMari @@ -3820,8 +3874,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1564049349000183809, // @DeboraMenezes22 1564101654496133120, - // @drbenedettirs - 1567868352432951298, + // @daiane_daiane18 + 1564304740850221058, + // @jaquepetrovik_ + 1565401487836061698, // @NMousquer 1569113308078170115, // @pitmagrin @@ -3862,10 +3918,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1586192715137667072, // @CombatPatriota 1586416104632721416, + // @LucasBackes14pr + 1586733661113794561, // @marcosfonsecapi 1587131451593658369, // @antoniodoidoofc 1588344483766321154, + // @munique_busson + 1588523690760732672, // @RONALDO90199231 1588685187629678592, // @gianninogueira2 @@ -3926,6 +3986,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1612520249735135232, // @depgilbertinho 1612589448256016385, + // @gersonzocchi1 + 1612930692932837377, // @WaldenorPereira 1613183981230366728, // @DiegoQuaqua @@ -3978,6 +4040,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1647755717158264835, // @wistongomess 1648046380106104833, + // @Alannleal_ + 1651297457374887936, // @CamilaGodoiSP 1653466465855471617, // @MatiasSamuka @@ -4000,14 +4064,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1668266165922144262, // @carrarajh23 1669323757700169729, - // @LuizGracianoMBL - 1674888775506313216, // @vanessarosajlle 1675490861440802817, + // @feschmittvet + 1676911232614293505, // @GiFreitas1982 1684192295459905536, - // @victormenezesrj - 1686531661305966592, // @fellipe1971 1690338078827728896, // @babatupinamba_ @@ -4036,12 +4098,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1710430261932908545, // @hebertcsgyn 1713936872370532352, + // @oengenheiroleo + 1714288040334831616, // @TitoBarichello 1715811403708166144, // @karisantospt 1716415063471316993, - // @manubarrossp - 1717241739302371328, // @moreiramissao 1721592893511483392, // @ThomazSJC @@ -4080,6 +4142,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1752837451175907328, // @rmpicoli 1757491174871363584, + // @prwellingtonst + 1758502910638329856, // @samuel_al_silva 1760365813456904192, // @RicardoAlv32716 @@ -4090,6 +4154,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1764371188489330689, // @JeffreyChiquini 1767171124629037056, + // @a_caroliveira + 1767282969708822528, // @Lenesilllva 1771346241449865216, // @rafaelsatiebr @@ -4102,6 +4168,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1776773968726298625, // @Ap52467Jovelino 1777297081134178304, + // @BrunoAlvesSFFS + 1777373652708683786, // @fabiocarneirojp 1777538059891777537, // @MarcoRo43166598 @@ -4116,6 +4184,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1779345802629914624, // @carlosiranrs 1779913435825795072, + // @BiaCoimbraSP + 1780537023859761153, // @CamillaGonda 1781485235047174144, // @brunnomattospt @@ -4138,8 +4208,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1793428861348167680, // @akarinaclaro 1793823661369044992, - // @RicardoSeneseUP - 1795906221662208000, // @nataliademesmao 1797413322385485825, // @Soldado_Sampaio @@ -4190,8 +4258,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1827352420541763584, // @MarcusLopesPsol 1827593492387745792, - // @drjoaomota832 - 1829618986171899905, // @matheussimoespr 1836487149769596928, // @AGoldbach60024 @@ -4204,6 +4270,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1845277420854837249, // @ProfMarcio58657 1845590017445400576, + // @henriquemgjf + 1845839091515981833, // @GoulartVla23085 1846568278950346752, // @DragUrbana @@ -4222,12 +4290,18 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1856816774517268480, // @Airtonjose26 1857173636689362944, + // @rubensnascimepr + 1857825522765766656, // @CristianeN74380 1858882502707810304, // @depcabomacielam 1859357204278542336, + // @LusFelipeV68715 + 1861380355875373056, // @denistaveiradn 1865917852424798208, + // @joaopaulo_tprs + 1870514350164705280, // @Pcbcastelo 1872641473021087744, // @saulofreitas22 @@ -4254,12 +4328,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1884620263549001728, // @yasminsarrafsp 1884758507272216576, + // @ravengaar + 1885085802645880832, // @ameliocayresdep 1887508211445702656, // @FelipeVasquesce 1889410539639517184, // @celprincipebr 1890456925147734016, + // @Herculano1011 + 1892340577796313092, // @julianafideliis 1896723491371581440, // @gualbertoap @@ -4278,6 +4356,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1899875749907357697, // @caixeta_oficial 1900205421123817472, + // @falaCosenza + 1900729566542655488, // @RafaMinatoSP 1901680009430892544, // @MacAntonioRJ @@ -4302,12 +4382,22 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1916570073663475712, // @aparecidobian 1920469901136764928, + // @ericafmissao26 + 1922067675074727936, + // @daniele_carva + 1925558068206518272, + // @MuriloM50730958 + 1928141892065316864, // @paulomeloparana 1929325555234803712, // @manoela__peres 1932149081620484096, + // @souadibelias + 1932784165545533440, // @leo_grandini13 1932924340917710848, + // @Camargo1Amanda + 1937005496398962688, // @prof_elson_sc 1937980838299504644, // @marcioalvinosp @@ -4370,6 +4460,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1986993348008419328, // @GustavoHenryR 1987721882838462464, + // @fig35508 + 1988116025670594563, // @DaClaudio33805 1988208205575712772, // @RenatoBolsonar0 @@ -4396,14 +4488,20 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2006965910309896192, // @alvarenga35289 2007441310102523904, + // @RogerioChimi + 2010017755194589185, // @AntoniadeJessp 2010525266565857280, + // @MichelPadao + 2010532868271816705, // @profraydf 2010751427623497728, // @edinhosouzaaa 2011825328793296896, // @Efreu_Quintana 2013273184640860160, + // @meunomeejhonebr + 2013398295859609601, // @brenomacedopi 2013660202873032707, // @esthermoraessp @@ -4412,22 +4510,36 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2015964533689323520, // @carolontiveros0 2016185817392099336, + // @anderalvesgo + 2016851339624538112, + // @Landovoigttt + 2017294158558339072, // @ingridcardososp 2017575066381238272, // @sicchar1389 2018445464908042240, + // @isabeldesouza01 + 2020307564903165953, // @vanesckaessusp 2021065723573829633, // @leninhavalente_ 2021229327594000386, + // @tamirisgodoysp + 2021276913012969472, // @MatheusCambuiBa 2021559145728454656, - // @Jordambritosc - 2022354011035222016, + // @ruandutrace + 2021639549625991168, + // @henriqueduraesd + 2021718062261506049, + // @sophiafechinece + 2022429125399478273, // @NetoFeitos68916 2026461466728292352, // @catarinanevespb 2026549782278529024, + // @jessicathis_mbl + 2026754488128909313, // @GreguyLoooban 2027415714953396224, // @AraceliLemosOF @@ -4438,10 +4550,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2029985728059559938, // @helen_vitaRJ 2033587304795963392, + // @CidaCarvalhoSP + 2034086382705274880, // @EdneyBatalha 2034242236678848513, - // @Brunodiasmissao - 2035542695536648192, // @Fabio_x86 2036151026986696704, // @MarcaoVivacqua @@ -4454,16 +4566,28 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2036894494973440000, // @isabeldesouzasp 2036944910713081856, + // @beatrizoPR + 2037168439836512257, + // @RicardoBec31758 + 2037169927618961408, // @stellabragasp 2040479116068081664, // @maurodeAL1930 2041451935685906432, // @DepDrFlavio 2041621533454450688, + // @oenzozibellini + 2041763427488608256, + // @Grazypasqualeto + 2042058510536482816, + // @CrisNavarroMIDH + 2043108953177899008, // @vanessacfortes 2043516611831701505, // @owilsonmartins 2043714770591735809, + // @CarolineSa5644 + 2043839578394533890, // @viniciusdiaspi 2045135821695574016, // @CMolinariBR @@ -4476,20 +4600,22 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2047324463612485632, // @JoaoPaulo_2026 2047395841715920897, - // @RitaDamore11 - 2047488478372368384, // @VasconcellosCel 2048607450593468416, // @PablodoMST 2048698490906169344, // @CGasparin11011 2048961988999454720, + // @livianoronha_PA + 2049127918794670080, // @MarcioRezendeRJ 2049539394701262848, // @bia_pedagoga 2049646600088117249, // @oalanmansurrj 2049927296442658816, + // @14DanielAguiar + 2050843854325067777, // @jaimeverruckms 2052755712644730887, // @Aminjhannouche @@ -4502,12 +4628,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2054619154762674176, // @onenencoelho 2054729517864771584, - // @KeremhadassaMG - 2054936868026777600, // @JulianaBrizola 2054954070083813376, // @mahmoudamer_rs 2054974955603861504, + // @DelEduardoK + 2055437449740873728, // @glaucelima12 2056410704031121408, // @edmartresoitao @@ -4516,6 +4642,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2056815030935416832, // @RodolfoFiorucci 2057456621865881600, + // @brenobarcelos_ + 2057550738427887616, // @PL22Al 2057808888297062400, // @ProfNelsiWelter @@ -4540,8 +4668,6 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2061280727711289345, // @DrCrisVeloso 2061889166779019264, - // @NilsonVicentisc - 2062522303892611072, // @DaversonMatos 2062692988606717952, // @jmonteirosc @@ -4550,12 +4676,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2062867137497059328, // @profjoaohs 2063783489275600896, + // @viniciusrodsp + 2064052780407336960, // @aBiaAlcantara 2064088166588440577, // @DelmaPSOL 2064731236023541760, // @pr_itamar_paim 2064819106545565696, + // @MalluCortes + 2065526840198832128, // @FelipeGambaroP 2066731525400330240, // @drcassiohprado @@ -4604,6 +4734,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2077757274961911808, // @LarianeTellMend 2077832017736011777, + // @DepCarlaMachado + 2078177903271956480, // @Nayladasilva0 2078543961333850112, // @drbrenoaraujo @@ -4660,6 +4792,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2082567723423342592, // @Clarianabr 2082581908005806081, + // @karimeefayad + 2082897339065233409, // @barbararesende0 2082924803145551872, // @guihenriquesc @@ -4690,6 +4824,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2084724540194643968, // @CarmemOliverofc 2084756136004083712, + // @PauloLealMissao + 2085083851395620864, // @panayotisdolula 2085133776741437441, // @diegojejees @@ -4704,12 +4840,24 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2086812805034815488, // @roneymariachi 2086873322587881473, + // @RaphaelaSpadi + 2087175409540497408, + // @capitaocesario + 2087749550719086592, // @DarbideJesusrr 2088266727519862784, // @tatibarrapa 2088334004445487104, // @Mariguedes1406 2088832528383639552, + // @PedroAbib15 + 2089024802669334528, + // @brunoscheid222 + 2089064009643220993, + // @SimonePimehb + 2090082297370312704, + // @cabodaciolo33 + 2091177221310345216, ]) }); @@ -4897,7 +5045,7 @@ mod tests { #[test] fn hardcoded_list_is_non_empty() { assert!(!BRAZIL_2026_ELECTION_USER_IDS.is_empty()); - assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2315); + assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2388); } #[test] diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index ddccebc3..cd728716 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -162,7 +162,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 { @@ -187,9 +187,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 { @@ -210,6 +210,7 @@ impl CandidateHelpers for PostCandidate { recon_gap_above: c.recon_gap_above, }), reward_rerank_slot_prob: None, + page_decode_slot_prob: None, reranker_head_tag: self.reranker_head_tag, } } diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index 80917b59..713dda5f 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-28T20:07:44Z +// mirrored from config feature-switch defaults; last sync 2026-08-31T16:17:17Z use xai_feature_switches::param; param!( @@ -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, 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 a0563925..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 { @@ -251,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/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/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 262a700f..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, @@ -274,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, @@ -295,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 3e446130..99a95c90 100644 --- a/phoenix/crates/common/xai-recsys/src/util.rs +++ b/phoenix/crates/common/xai-recsys/src/util.rs @@ -358,7 +358,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]; @@ -435,6 +435,10 @@ impl InputBuffer { [base_idx..base_idx + search_query_embedding_dim] .copy_from_slice(&candidate_set.search_query_embedding); } + } else { + log::error!( + "search_query_embedding dim {provided_dim} != model {search_query_embedding_dim}; leaving zeros" + ); } } @@ -1550,7 +1554,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, diff --git a/phoenix/crates/serving/xai-recsys-engine/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index 650a11ec..74bbfda7 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -1877,7 +1877,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(|| { 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..11082080 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,128 @@ 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_sync( &self, post_ids: Vec, emb_dim: usize, candidate_seq_len: usize, ) -> Result> { + use rayon::prelude::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + 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 missing = AtomicUsize::new(0); + + results + .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"]) @@ -211,6 +257,15 @@ impl MmEmbeddingsClient { FETCH_MM_EMBEDDINGS_DURATION.observe(start_time.elapsed().as_secs_f64()); 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 +283,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 +314,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 +392,33 @@ 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); + } + #[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 e9100296..17e6a5c2 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -112,6 +112,24 @@ message PredictNextActionsRequest { repeated uint32 requestedContinuousActionIndices = 17; repeated ConvAssetIds conv_asset_ids = 18; + + PageDecodeParams pageDecode = 19; +} + +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 { @@ -200,6 +218,8 @@ message CandidateDistributionSet { string checkpointPath = 3; string intermediateTensorsPath = 4; + + PageDecodeResult pageDecodeResult = 5; } message ScoredCandidates { @@ -423,57 +443,18 @@ enum ActionName { 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; + 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; P_OPEN_LINK_P90 = 253; - PLACE_HOLDER_254 = 254; + ADS_MMP_CLICK = 254; PLACE_HOLDER_255 = 255; } @@ -1237,6 +1218,7 @@ message ScoreInfo { SlateContext slateContext = 4; optional double rewardRerankSlotProb = 5; optional uint32 rerankerHeadTag = 6; + optional double pageDecodeSlotProb = 7; } message SlateContext { diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index e9100296..17e6a5c2 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -112,6 +112,24 @@ message PredictNextActionsRequest { repeated uint32 requestedContinuousActionIndices = 17; repeated ConvAssetIds conv_asset_ids = 18; + + PageDecodeParams pageDecode = 19; +} + +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 { @@ -200,6 +218,8 @@ message CandidateDistributionSet { string checkpointPath = 3; string intermediateTensorsPath = 4; + + PageDecodeResult pageDecodeResult = 5; } message ScoredCandidates { @@ -423,57 +443,18 @@ enum ActionName { 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; + 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; P_OPEN_LINK_P90 = 253; - PLACE_HOLDER_254 = 254; + ADS_MMP_CLICK = 254; PLACE_HOLDER_255 = 255; } @@ -1237,6 +1218,7 @@ message ScoreInfo { SlateContext slateContext = 4; optional double rewardRerankSlotProb = 5; optional uint32 rerankerHeadTag = 6; + optional double pageDecodeSlotProb = 7; } message SlateContext { 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..5bf88e19 --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import json +import pathlib +import time + +TREE_DEK_NAME = "_DEK" +TREE_DEK_CLAIM_NAME = "_DEK.claim" +PUBLISH_TIMEOUT_SECS = 120.0 + + +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]: + import xai_kms + + dek_path = path / TREE_DEK_NAME + claim = path / TREE_DEK_CLAIM_NAME + + deadline = time.monotonic() + PUBLISH_TIMEOUT_SECS + 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) + 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" + ) + time.sleep(0.1) + raw = bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, dek_path)) + return raw, entry["wrapped"], entry.get("context") or {} 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..ef63e751 --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 X.AI Corp. +import base64 +import pathlib + +from xai_checkpointing.dek import publish_tree_dek + + +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 = publish_tree_dek(path, kms_client) + return _envelope_spec( + path, base64.b64encode(raw).decode(), wrapped, encryption_chunk_size, context + ) diff --git a/phoenix/xrex/configs/data_feeds.py b/phoenix/xrex/configs/data_feeds.py index 917ef4be..1eeb95af 100644 --- a/phoenix/xrex/configs/data_feeds.py +++ b/phoenix/xrex/configs/data_feeds.py @@ -292,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, ) @@ -326,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_two_tower.py b/phoenix/xrex/configs/xrecsys_two_tower.py index 0b5ce8d4..b35329d7 100644 --- a/phoenix/xrex/configs/xrecsys_two_tower.py +++ b/phoenix/xrex/configs/xrecsys_two_tower.py @@ -333,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 = { 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/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/utils/checkpointing.py b/phoenix/xrex/utils/checkpointing.py index cd287d22..12ad67a3 100644 --- a/phoenix/xrex/utils/checkpointing.py +++ b/phoenix/xrex/utils/checkpointing.py @@ -233,10 +233,7 @@ def save_checkpoint( 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): diff --git a/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala b/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala index 5b9f70b4..c1323117 100644 --- a/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala +++ b/simclusters/simclusters_v2/scalding/embedding/EntityToSimClustersEmbeddingsJob.scala @@ -5,10 +5,7 @@ import com.twitter.recos.entities.thriftscala.Entity import com.twitter.recos.entities.thriftscala.Hashtag import com.twitter.recos.entities.thriftscala.SemanticCoreEntity import com.twitter.scalding._ -import com.twitter.scalding_internal.dalv2.DAL import com.twitter.scalding_internal.dalv2.DALWrite._ -import com.twitter.scalding_internal.dalv2.remote_access.ExplicitLocation -import com.twitter.scalding_internal.dalv2.remote_access.ProcAtla import com.twitter.scalding_internal.multiformat.format.keyval.KeyVal import com.twitter.simclusters_v2.common.ModelVersions import com.twitter.simclusters_v2.common.SimClustersEmbedding @@ -215,20 +212,8 @@ trait EntityToSimClustersEmbeddingApp extends ScheduledExecutionApp { val simClustersEmbedding = jobConfig.modelVersion match { case ModelVersion.Model20m145k2020 => - val interestedIn2020WithFallback = - SimclustersV2InterestedIn20M145K2020ScalaDataset.copy(fallbackPath = Some("viewfs://hadoop-nn.example.invalid/user/cassowary/manhattan_sequence_files/" + - s"simclusters_v2_interested_in_20M_145K_2020/_TMP_RECOVERY/${dateRange.start.timestamp}")) - val simClustersSource2020 = DAL - .readMostRecentSnapshot( - interestedIn2020WithFallback, - dateRange.prepend(Days(28)(timeZone)) - ) - .withRemoteReadPolicy(ExplicitLocation(ProcAtla)) - .toTypedPipe - .map { - case KeyVal(userId, clustersUserIsInterestedIn) => - (userId, clustersUserIsInterestedIn) - } + val simClustersSource2020 = + InterestedInSources.simClustersInterestedIn2020Source(dateRange, timeZone) computeEmbeddings( simClustersSource2020, normalizedUserEntityMatrix, 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/config.rs b/visibility-filtering/config.rs index c67f1ea1..be32bb26 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -8,9 +8,16 @@ pub const ENV_FALLBACK_CACHE_SERVE_STALE_ENABLED: &str = "VF_FALLBACK_CACHE_SERV pub const ENV_FALLBACK_CACHE_POPULATE_ENABLED: &str = "VF_FALLBACK_CACHE_POPULATE_ENABLED"; pub const ENV_CACHE_WARM_SAMPLE_PCT: &str = "VF_CACHE_WARM_SAMPLE_PCT"; 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"; +pub 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 gizmoduck_client_id() -> String { resolve_gizmoduck_client_id( std::env::var(ENV_GIZMODUCK_CLIENT_ID).ok().as_deref(), diff --git a/visibility-filtering/filter.rs b/visibility-filtering/filter.rs index 18aeaa95..5be5b357 100644 --- a/visibility-filtering/filter.rs +++ b/visibility-filtering/filter.rs @@ -121,7 +121,7 @@ impl FilterTweets { } #[cfg(test)] -mod tests { +pub(crate) mod test_support { use super::*; use crate::clients::socialgraph_client::MockSocialgraphClient; use crate::safety_label_source::lookup::{ManhattanLookup, RemoteSource, TwemcacheLookup}; @@ -169,10 +169,14 @@ 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 twemcache = Arc::new(FakeTwemcache); let manhattan = Arc::new(FakeManhattan); @@ -191,6 +195,12 @@ mod tests { Policies::new(), ) } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::test_support::filter_tweets; fn candidate(tweet_id: u64, author_id: Option) -> RawCandidate { RawCandidate { diff --git a/visibility-filtering/filter_tweets.rs b/visibility-filtering/filter_tweets.rs index 1641ff31..79aa575a 100644 --- a/visibility-filtering/filter_tweets.rs +++ b/visibility-filtering/filter_tweets.rs @@ -33,6 +33,8 @@ impl FilterTweetsEndpoint { 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, @@ -70,7 +72,7 @@ impl FilterTweetsEndpoint { let response = self .filter_tweets .run(FilterRequest { - viewer_id: req.viewer_id, + viewer_id, country_code: req.country_code, safety_level, candidates, @@ -109,6 +111,10 @@ impl FilterTweetsEndpoint { } } +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), @@ -145,6 +151,45 @@ mod tests { } } + 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 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 allow_maps_to_proto_result() { let result = to_visibility_result(outcome(7, VfAction::Allow)); diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index df8bc331..4b967ea1 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -6,6 +6,7 @@ pub(crate) mod filter_tweets; pub(crate) mod get_safety_labels; pub mod hydration; pub mod models; +pub mod params; pub(crate) mod reference_compare; pub mod rules; pub mod safety_label_source; diff --git a/visibility-filtering/main.rs b/visibility-filtering/main.rs index f404a789..4d6c4c99 100644 --- a/visibility-filtering/main.rs +++ b/visibility-filtering/main.rs @@ -25,6 +25,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) diff --git a/visibility-filtering/params.rs b/visibility-filtering/params.rs new file mode 100644 index 00000000..af4dcf3d --- /dev/null +++ b/visibility-filtering/params.rs @@ -0,0 +1,216 @@ +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"; + +pub const SCALA_NSFW_GATING_FILE: &str = "country_specific_nsfw_content_gating.yml"; +pub 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 struct NsfwGatingCountries { + countries: ArcSwap>, +} + +impl NsfwGatingCountries { + pub fn new() -> 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) + } + + 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); + } + }); + } +} + +impl Default for NsfwGatingCountries { + fn default() -> Self { + Self::new() + } +} + +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::new(); + 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::new(); + 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 index 31ae39d4..86040c0d 100644 --- a/visibility-filtering/reference_compare.rs +++ b/visibility-filtering/reference_compare.rs @@ -338,7 +338,7 @@ impl ReferenceCompareHarness { harness } - pub(crate) fn begin_compare( + pub(crate) fn begin_compare( self: &Arc, viewer_id: Option, country_code: Option, @@ -370,13 +370,17 @@ impl ReferenceCompareHarness { let (reference_outcome, verdicts) = futures::future::join(tokio::time::timeout(REFERENCE_TIMEOUT, reference_fut), rx) .await; - let reference_results = match reference_outcome { - Ok(results) => results, - Err(_) => { - harness.incr(ERROR, &[("kind", "timeout")]); - return; - } - }; + 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, @@ -427,6 +431,8 @@ mod tests { 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 @@ -625,7 +631,7 @@ mod tests { } } - const ID: u64 = 1_000_000_000_000_000_000; + const ID: u64 = 1_000_000_000_000_000_000; #[test] fn identical_pairs_group_into_one_diffs_entry() { @@ -728,8 +734,19 @@ mod tests { safety_level: ReferenceSafetyLevel, for_user_id: u64, context: Option, - ) -> HashMap>> { - let results = tweet_ids.iter().map(|&id| (id, Ok(None))).collect(); + ) -> 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, diff --git a/visibility-filtering/rules/context.rs b/visibility-filtering/rules/context.rs index 5a5a0135..b2c60141 100644 --- a/visibility-filtering/rules/context.rs +++ b/visibility-filtering/rules/context.rs @@ -22,177 +22,268 @@ impl<'a> RuleContext<'a> { } } + #[inline] pub fn safety_level(&self) -> SafetyLevel { self.safety_level } - pub fn viewer_is_logged_out(&self) -> bool { - self.viewer.viewer_is_logged_out() + #[inline] + pub fn viewer(&self) -> ViewerPredicates<'_> { + ViewerPredicates { ctx: self } } - pub fn viewer_is_underage(&self) -> bool { - self.viewer.viewer_is_underage() + #[inline] + pub fn tweet(&self) -> TweetPredicates<'_> { + TweetPredicates { ctx: self } } - pub fn viewer_has_no_stated_age(&self) -> bool { - self.viewer.viewer_has_no_stated_age() + #[inline] + pub fn author(&self) -> AuthorPredicates<'_> { + AuthorPredicates { ctx: self } } - pub fn viewer_allows_sensitive_media(&self) -> bool { - self.viewer.allows_sensitive_media + #[inline] + pub fn takedown(&self) -> TakedownPredicates<'_> { + TakedownPredicates { ctx: self } + } +} + +#[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() } - pub fn viewer_country_in(&self, countries: &[&str]) -> bool { - self.viewer + #[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.viewer.country_code.as_deref()) - .is_some_and(|c| countries.contains(&c)) + .or(self.ctx.viewer.country_code.as_deref()) } - pub fn is_author_viewer(&self) -> bool { - self.candidate.is_author_viewer(self.viewer.viewer) + #[inline] + pub fn is_author(&self) -> bool { + self.ctx.candidate.is_author_viewer(self.ctx.viewer.viewer) } - pub fn viewer_follows_author(&self) -> bool { - self.candidate.viewer_follows_author() + #[inline] + pub fn follows_author(&self) -> bool { + self.ctx.candidate.viewer_follows_author() } - pub fn viewer_blocks_author(&self) -> bool { - self.candidate.relationship.viewer_blocks_author + #[inline] + pub fn blocks_author(&self) -> bool { + self.ctx.candidate.relationship.viewer_blocks_author } - pub fn viewer_mutes_author(&self) -> bool { - self.candidate.relationship.viewer_mutes_author + #[inline] + pub fn mutes_author(&self) -> bool { + self.ctx.candidate.relationship.viewer_mutes_author } - pub fn viewer_mutes_retweets_from_author(&self) -> bool { - self.candidate + #[inline] + pub fn mutes_retweets_from_author(&self) -> bool { + self.ctx + .candidate .relationship .viewer_mutes_retweets_from_author } - pub fn has_tweet_safety_label(&self, label: SafetyLabelType) -> bool { - self.candidate.has_safety_label(label) + #[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.candidate.is_retweet() + self.ctx.candidate.is_retweet() } - pub fn is_stale_tweet(&self) -> bool { - self.candidate.is_stale_tweet() + #[inline] + pub fn is_stale(&self) -> bool { + self.ctx.candidate.is_stale_tweet() } + #[inline] pub fn is_nullcast(&self) -> bool { - self.candidate.is_nullcast() + self.ctx.candidate.is_nullcast() } + #[inline] pub fn is_community_tweet(&self) -> bool { - self.candidate.is_community_tweet() + self.ctx.candidate.is_community_tweet() } + #[inline] pub fn has_media(&self) -> bool { - self.candidate.has_media() + self.ctx.candidate.has_media() } + #[inline] pub fn has_dmca_media(&self) -> bool { - self.candidate.has_dmca_media() + self.ctx.candidate.has_dmca_media() } + #[inline] pub fn is_nsfw_flagged(&self) -> bool { - self.candidate.is_nsfw_flagged() + self.ctx.candidate.is_nsfw_flagged() } - pub fn has_tweet_nsfw_user_flag(&self) -> bool { - self.candidate.tweet_features.nsfw.user + #[inline] + pub fn has_nsfw_user_flag(&self) -> bool { + self.ctx.candidate.tweet_features.nsfw.user } - pub fn has_tweet_nsfw_admin_flag(&self) -> bool { - self.candidate.tweet_features.nsfw.admin + #[inline] + pub fn has_nsfw_admin_flag(&self) -> bool { + self.ctx.candidate.tweet_features.nsfw.admin } - pub fn legal_takedown_in_viewer_country(&self) -> bool { - self.takedown_in_viewer_country(legal_takedown_country) + #[inline] + pub fn is_exclusive(&self) -> bool { + self.ctx.candidate.exclusive_content.is_some() } +} - pub fn local_laws_takedown_in_viewer_country(&self) -> bool { - self.takedown_in_viewer_country(local_laws_takedown_country) - } +#[derive(Clone, Copy)] +pub struct AuthorPredicates<'a> { + ctx: &'a RuleContext<'a>, +} - fn takedown_in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool { - let Some(viewer_country) = &self.viewer.country_code else { - return false; - }; - self.candidate - .tweet_features - .takedown - .reasons - .iter() - .filter_map(extractor) - .any(|c| c.eq_ignore_ascii_case(viewer_country)) +impl AuthorPredicates<'_> { + #[inline] + pub fn is_suspended(&self) -> bool { + self.ctx.candidate.author_features.is_suspended } - pub fn media_restricted_in_viewer_country(&self) -> bool { - let country = self - .viewer - .country_code - .as_deref() - .unwrap_or(WORLDWIDE_COUNTRY_CODE); - let allow = &self.candidate.tweet_features.media.geo_allow_list; - let deny = &self.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)) + #[inline] + pub fn is_deactivated(&self) -> bool { + self.ctx.candidate.author_features.is_deactivated } - pub fn author_is_suspended(&self) -> bool { - self.candidate.author_features.is_suspended + #[inline] + pub fn is_erased(&self) -> bool { + self.ctx.candidate.author_features.is_erased } - pub fn author_is_deactivated(&self) -> bool { - self.candidate.author_features.is_deactivated + #[inline] + pub fn is_offboarded(&self) -> bool { + self.ctx.candidate.author_features.is_offboarded } - pub fn author_is_erased(&self) -> bool { - self.candidate.author_features.is_erased + #[inline] + pub fn is_protected(&self) -> bool { + self.ctx.candidate.author_features.is_protected } - pub fn author_is_offboarded(&self) -> bool { - self.candidate.author_features.is_offboarded + #[inline] + pub fn is_nsfw_user(&self) -> bool { + self.ctx.candidate.author_features.is_nsfw_user } - pub fn author_is_protected(&self) -> bool { - self.candidate.author_features.is_protected + #[inline] + pub fn is_nsfw_admin(&self) -> bool { + self.ctx.candidate.author_features.is_nsfw_admin } - pub fn author_is_nsfw_user(&self) -> bool { - self.candidate.author_features.is_nsfw_user + #[inline] + pub fn has_user_label(&self, label: LabelValue) -> bool { + self.ctx.candidate.author_has_user_label(label) } +} - pub fn author_is_nsfw_admin(&self) -> bool { - self.candidate.author_features.is_nsfw_admin - } +#[derive(Clone, Copy)] +pub struct TakedownPredicates<'a> { + ctx: &'a RuleContext<'a>, +} - pub fn author_has_user_label(&self, label: LabelValue) -> bool { - self.candidate.author_has_user_label(label) +impl TakedownPredicates<'_> { + #[inline] + pub fn legal_in_viewer_country(&self) -> bool { + self.in_viewer_country(legal_takedown_country) } - pub fn is_exclusive_tweet(&self) -> bool { - self.candidate.exclusive_content.is_some() + #[inline] + pub fn local_laws_in_viewer_country(&self) -> bool { + self.in_viewer_country(local_laws_takedown_country) } - pub fn viewer_is_conversation_author(&self) -> bool { - match (&self.candidate.exclusive_content, self.viewer.viewer_id()) { - (Some(exclusive), Some(viewer_id)) => viewer_id == exclusive.conversation_author_id, - _ => false, - } + #[inline] + fn in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool { + let Some(viewer_country) = &self.ctx.viewer.country_code else { + return false; + }; + self.ctx + .candidate + .tweet_features + .takedown + .reasons + .iter() + .filter_map(extractor) + .any(|c| c.eq_ignore_ascii_case(viewer_country)) } - pub fn viewer_super_follows_author(&self) -> bool { - self.candidate - .exclusive_content - .as_ref() - .is_some_and(|exclusive| exclusive.viewer_super_follows_author) + #[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)) } } 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/nsfw_age_gating.rs b/visibility-filtering/rules/nsfw_age_gating.rs index 1146c2b2..3c4ab4cf 100644 --- a/visibility-filtering/rules/nsfw_age_gating.rs +++ b/visibility-filtering/rules/nsfw_age_gating.rs @@ -1,29 +1,30 @@ use crate::models::{SafetyLabelType, VfAction}; +use crate::params::NsfwGatingCountries; use crate::rules::{Rule, RuleContext}; +use std::sync::Arc; 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 { - !context.is_author_viewer() - && context.has_media() - && (context.has_tweet_safety_label(SafetyLabelType::NSFW_HIGH_PRECISION) - || context.has_tweet_safety_label(SafetyLabelType::NSFW_HIGH_RECALL) - || (context.is_nsfw_flagged() && !context.is_retweet())) + 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 { - !context.is_author_viewer() - && context.has_media() - && context.has_tweet_safety_label(SafetyLabelType::GORE_AND_VIOLENCE_HIGH_PRECISION) + 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 { - !context.is_author_viewer() - && (context.has_tweet_safety_label(SafetyLabelType::NSFW_TEXT) - || context.has_tweet_safety_label(SafetyLabelType::NSFW_CARD_IMAGE)) + 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 { @@ -40,7 +41,7 @@ impl Rule for SensitiveViewerLoggedOutDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_is_logged_out() && sensitive_base_condition(context) { + if context.viewer().is_logged_out() && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) } else { VfAction::Allow @@ -56,7 +57,7 @@ impl Rule for SensitiveViewerUnderageDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_is_underage() && sensitive_base_condition(context) { + if context.viewer().is_underage() && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) } else { VfAction::Allow @@ -64,7 +65,15 @@ impl Rule for SensitiveViewerUnderageDropRule { } } -pub struct SensitiveViewerNoStatedAgeDropRule; +pub struct SensitiveViewerNoStatedAgeDropRule { + gating_countries: Arc, +} + +impl SensitiveViewerNoStatedAgeDropRule { + pub fn new(gating_countries: Arc) -> Self { + Self { gating_countries } + } +} impl Rule for SensitiveViewerNoStatedAgeDropRule { fn name(&self) -> &'static str { @@ -72,8 +81,11 @@ impl Rule for SensitiveViewerNoStatedAgeDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_has_no_stated_age() - && context.viewer_country_in(&NSFW_GATING_COUNTRIES) + if context.viewer().has_no_stated_age() + && context + .viewer() + .country() + .is_some_and(|c| self.gating_countries.contains(c)) && sensitive_base_condition(context) { VfAction::Drop(FilteredReason::ContainNsfwMedia) @@ -103,6 +115,10 @@ mod tests { candidate().with_label(label).with_media().build() } + fn no_stated_age_rule() -> SensitiveViewerNoStatedAgeDropRule { + SensitiveViewerNoStatedAgeDropRule::new(Arc::new(NsfwGatingCountries::new())) + } + fn nsfw_author_media_candidate() -> HydratedTweetCandidate { candidate() .with_media() @@ -185,7 +201,7 @@ mod tests { VfAction::Allow )); assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( + no_stated_age_rule().evaluate(&crate::rules::test_context( &gating_viewer(ViewerAge::Unknown), &c )), @@ -227,7 +243,7 @@ mod tests { 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( + no_stated_age_rule().evaluate(&crate::rules::test_context( &gating_viewer(ViewerAge::NotStated), &c )), @@ -243,7 +259,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Allow )); } @@ -297,7 +313,7 @@ mod tests { VfAction::Allow )); assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context( + no_stated_age_rule().evaluate(&crate::rules::test_context( &gating_viewer(ViewerAge::Unknown), &c )), @@ -348,7 +364,7 @@ mod tests { 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( + no_stated_age_rule().evaluate(&crate::rules::test_context( &gating_viewer(ViewerAge::NotStated), &c )), @@ -364,7 +380,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Allow )); } @@ -377,7 +393,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Allow )); } @@ -391,7 +407,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Allow )); } @@ -405,7 +421,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Drop(_) )); } @@ -419,7 +435,7 @@ mod tests { ..gating_viewer(ViewerAge::NotStated) }; assert!(matches!( - SensitiveViewerNoStatedAgeDropRule.evaluate(&crate::rules::test_context(&v, &c)), + no_stated_age_rule().evaluate(&crate::rules::test_context(&v, &c)), VfAction::Drop(_) )); } diff --git a/visibility-filtering/rules/nsfw_interstitial.rs b/visibility-filtering/rules/nsfw_interstitial.rs index 674de6fa..6438c02f 100644 --- a/visibility-filtering/rules/nsfw_interstitial.rs +++ b/visibility-filtering/rules/nsfw_interstitial.rs @@ -20,9 +20,9 @@ impl Rule for NsfwMediaInterstitialRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.has_tweet_safety_label(self.label) - && !context.is_author_viewer() - && !context.viewer_allows_sensitive_media() + if context.tweet().has_safety_label(self.label) + && !context.viewer().is_author() + && !context.viewer().allows_sensitive_media() { return VfAction::Interstitial(FilteredReason::ContainNsfwMedia); } @@ -55,10 +55,10 @@ impl Rule for NsfwAuthorInterstitialRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.is_nsfw_flagged() - && context.has_media() - && !context.is_author_viewer() - && !context.viewer_allows_sensitive_media() + if context.tweet().is_nsfw_flagged() + && context.tweet().has_media() + && !context.viewer().is_author() + && !context.viewer().allows_sensitive_media() { return VfAction::Interstitial(FilteredReason::ContainNsfwMedia); } diff --git a/visibility-filtering/rules/nullcast_rule.rs b/visibility-filtering/rules/nullcast_rule.rs index db4e22b7..c677ab66 100644 --- a/visibility-filtering/rules/nullcast_rule.rs +++ b/visibility-filtering/rules/nullcast_rule.rs @@ -10,7 +10,10 @@ impl Rule for NullcastedTweetDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.is_nullcast() && !context.is_retweet() && !context.is_community_tweet() { + if context.tweet().is_nullcast() + && !context.tweet().is_retweet() + && !context.tweet().is_community_tweet() + { return VfAction::Drop(FilteredReason::TweetIsNullcast); } VfAction::Allow diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index acf6f96f..5c7d204b 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -1,4 +1,5 @@ use crate::models::{HydratedTweetCandidate, VfAction, ViewerFeatures}; +use crate::params::NsfwGatingCountries; use crate::rules::nsfw_age_gating::{ SensitiveViewerLoggedOutDropRule, SensitiveViewerNoStatedAgeDropRule, SensitiveViewerUnderageDropRule, @@ -20,6 +21,7 @@ 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 std::sync::Arc; use xai_visibility_filtering::models::FilteredReason; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -47,10 +49,14 @@ pub struct Policies { impl Policies { pub fn new() -> Self { + Self::with_nsfw_gating_countries(Arc::new(NsfwGatingCountries::new())) + } + + pub fn with_nsfw_gating_countries(gating_countries: Arc) -> Self { Self { filter_all: vec![Box::new(FilterAllRule)], - timeline_home: timeline_home_policy(), - timeline_home_recommendations: timeline_home_recommendations_policy(), + timeline_home: timeline_home_policy(&gating_countries), + timeline_home_recommendations: timeline_home_recommendations_policy(&gating_countries), } } @@ -103,7 +109,7 @@ impl Rule for FilterAllRule { } } -fn base_home_rules() -> Vec> { +fn base_home_rules(gating_countries: &Arc) -> Vec> { vec![ Box::new(author::SUSPENDED_AUTHOR_DROP), Box::new(author::DEACTIVATED_AUTHOR_DROP), @@ -127,7 +133,9 @@ fn base_home_rules() -> Vec> { Box::new(DropLocalLawsTakendownPostRule), Box::new(SensitiveViewerLoggedOutDropRule), Box::new(SensitiveViewerUnderageDropRule), - Box::new(SensitiveViewerNoStatedAgeDropRule), + Box::new(SensitiveViewerNoStatedAgeDropRule::new(Arc::clone( + gating_countries, + ))), Box::new(DropExclusiveTweetContentRule), Box::new(NSFW_HIGH_PRECISION_INTERSTITIAL), Box::new(GORE_AND_VIOLENCE_INTERSTITIAL), @@ -136,12 +144,14 @@ fn base_home_rules() -> Vec> { ] } -fn timeline_home_policy() -> Vec> { - base_home_rules() +fn timeline_home_policy(gating_countries: &Arc) -> Vec> { + base_home_rules(gating_countries) } -fn timeline_home_recommendations_policy() -> Vec> { - let mut rules = base_home_rules(); +fn timeline_home_recommendations_policy( + gating_countries: &Arc, +) -> Vec> { + let mut rules = base_home_rules(gating_countries); let oon_drops: Vec> = vec![ Box::new(DropTweetsWithDmcaMediaRule), Box::new(DropTweetsWithGeoRestrictedMediaRule), @@ -225,6 +235,44 @@ mod tests { )); } + #[test] + fn refreshed_config_country_reaches_the_wired_rule() { + let gating_countries = Arc::new(NsfwGatingCountries::new()); + let policies = Policies::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_age: crate::models::ViewerAge::NotStated, + country_code: Some("us".into()), + ..viewer(VIEWER_ID) + }; + + let verdict = policies.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" +"#, + ) + .unwrap(), + ); + let verdict = policies.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); + assert!(matches!(verdict.action, VfAction::Drop(_))); + assert_eq!( + verdict.decided_by, + Some("SensitiveViewerNoStatedAgeDropRule") + ); + } + #[test] fn filter_all_rule_drops_even_self_view() { let candidate = candidate().build(); diff --git a/visibility-filtering/rules/socialgraph_rules.rs b/visibility-filtering/rules/socialgraph_rules.rs index c9b4707e..79815221 100644 --- a/visibility-filtering/rules/socialgraph_rules.rs +++ b/visibility-filtering/rules/socialgraph_rules.rs @@ -10,10 +10,10 @@ impl Rule for ViewerBlocksAuthorRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_is_logged_out() { + if context.viewer().is_logged_out() { return VfAction::Allow; } - if context.viewer_blocks_author() { + if context.viewer().blocks_author() { return VfAction::Drop(FilteredReason::AuthorBlockViewer); } VfAction::Allow @@ -28,10 +28,10 @@ impl Rule for MutedRetweetsRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_is_logged_out() { + if context.viewer().is_logged_out() { return VfAction::Allow; } - if context.is_retweet() && context.viewer_mutes_retweets_from_author() { + if context.tweet().is_retweet() && context.viewer().mutes_retweets_from_author() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -46,10 +46,10 @@ impl Rule for ViewerMutesAuthorRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer_is_logged_out() { + if context.viewer().is_logged_out() { return VfAction::Allow; } - if context.viewer_mutes_author() { + if context.viewer().mutes_author() { return VfAction::Drop(FilteredReason::ViewerMutesAuthor); } VfAction::Allow @@ -64,23 +64,23 @@ impl Rule for DropExclusiveTweetContentRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.is_exclusive_tweet() { + if !context.tweet().is_exclusive() { return VfAction::Allow; } - if context.viewer_is_logged_out() { + if context.viewer().is_logged_out() { return VfAction::Drop(FilteredReason::ExclusiveTweet); } - if context.viewer_is_conversation_author() { + if context.viewer().is_conversation_author() { return VfAction::Allow; } - if context.viewer_super_follows_author() { + if context.viewer().super_follows_author() { return VfAction::Allow; } - if !context.is_retweet() && context.is_author_viewer() { + if !context.tweet().is_retweet() && context.viewer().is_author() { return VfAction::Allow; } diff --git a/visibility-filtering/rules/tes_rules.rs b/visibility-filtering/rules/tes_rules.rs index b65e0ed7..35df980f 100644 --- a/visibility-filtering/rules/tes_rules.rs +++ b/visibility-filtering/rules/tes_rules.rs @@ -10,7 +10,7 @@ impl Rule for DropStaleTweetsRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.is_stale_tweet() && !context.is_retweet() { + if context.tweet().is_stale() && !context.tweet().is_retweet() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -25,7 +25,7 @@ impl Rule for DropLegalTakendownPostRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.is_author_viewer() && context.legal_takedown_in_viewer_country() { + if !context.viewer().is_author() && context.takedown().legal_in_viewer_country() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -40,7 +40,7 @@ impl Rule for DropLocalLawsTakendownPostRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.is_author_viewer() && context.local_laws_takedown_in_viewer_country() { + if !context.viewer().is_author() && context.takedown().local_laws_in_viewer_country() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -55,7 +55,7 @@ impl Rule for DropTweetsWithGeoRestrictedMediaRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.media_restricted_in_viewer_country() { + if context.takedown().media_restricted_in_viewer_country() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow @@ -70,7 +70,7 @@ impl Rule for DropTweetsWithDmcaMediaRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.has_dmca_media() { + if context.tweet().has_dmca_media() { return VfAction::Drop(FilteredReason::UnspecifiedReason); } VfAction::Allow diff --git a/visibility-filtering/rules/tweet_flag_rules.rs b/visibility-filtering/rules/tweet_flag_rules.rs index 18dac8f7..06d6d233 100644 --- a/visibility-filtering/rules/tweet_flag_rules.rs +++ b/visibility-filtering/rules/tweet_flag_rules.rs @@ -34,12 +34,12 @@ impl Rule for TweetFlagDropRule { pub const TWEET_NSFW_USER_DROP: TweetFlagDropRule = TweetFlagDropRule::new( "TweetNsfwUserDropRule", - |context| context.has_tweet_nsfw_user_flag(), + |context| context.tweet().has_nsfw_user_flag(), FilteredReason::ContainNsfwMedia, ); pub const TWEET_NSFW_ADMIN_DROP: TweetFlagDropRule = TweetFlagDropRule::new( "TweetNsfwAdminDropRule", - |context| context.has_tweet_nsfw_admin_flag(), + |context| context.tweet().has_nsfw_admin_flag(), FilteredReason::ContainNsfwMedia, ); diff --git a/visibility-filtering/rules/tweet_label_drops.rs b/visibility-filtering/rules/tweet_label_drops.rs index a915760f..32bcab93 100644 --- a/visibility-filtering/rules/tweet_label_drops.rs +++ b/visibility-filtering/rules/tweet_label_drops.rs @@ -34,10 +34,10 @@ impl Rule for SafetyLabelDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.has_tweet_safety_label(self.label) { + if !context.tweet().has_safety_label(self.label) { return VfAction::Allow; } - if self.exempt_author && context.is_author_viewer() { + if self.exempt_author && context.viewer().is_author() { return VfAction::Allow; } VfAction::Drop(self.reason.clone()) diff --git a/visibility-filtering/rules/user_label_drops.rs b/visibility-filtering/rules/user_label_drops.rs index 09b6fdd1..e72c82e7 100644 --- a/visibility-filtering/rules/user_label_drops.rs +++ b/visibility-filtering/rules/user_label_drops.rs @@ -33,15 +33,15 @@ impl Rule for UserSafetyLabelDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.is_author_viewer() { + if context.viewer().is_author() { return VfAction::Allow; } - if !context.author_has_user_label(self.label) { + if !context.author().has_user_label(self.label) { return VfAction::Allow; } if self.require_non_follower - && !context.viewer_is_logged_out() - && context.viewer_follows_author() + && !context.viewer().is_logged_out() + && context.viewer().follows_author() { return VfAction::Allow; } diff --git a/visibility-filtering/rules/user_rules.rs b/visibility-filtering/rules/user_rules.rs index b00f6d9b..d9ca8a87 100644 --- a/visibility-filtering/rules/user_rules.rs +++ b/visibility-filtering/rules/user_rules.rs @@ -25,7 +25,7 @@ impl Rule for AuthorFlagDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if (self.flag)(context) && !context.is_author_viewer() { + if (self.flag)(context) && !context.viewer().is_author() { return VfAction::Drop(self.reason.clone()); } VfAction::Allow @@ -34,32 +34,32 @@ impl Rule for AuthorFlagDropRule { pub const SUSPENDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "SuspendedAuthorRule", - |context| context.author_is_suspended(), + |context| context.author().is_suspended(), FilteredReason::AuthorIsSuspended, ); pub const DEACTIVATED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DeactivatedAuthorRule", - |context| context.author_is_deactivated(), + |context| context.author().is_deactivated(), FilteredReason::AuthorIsDeactivated, ); pub const ERASED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "ErasedAuthorRule", - |context| context.author_is_erased(), + |context| context.author().is_erased(), FilteredReason::AuthorAccountIsInactive, ); pub const OFFBOARDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "OffboardedAuthorRule", - |context| context.author_is_offboarded(), + |context| context.author().is_offboarded(), FilteredReason::AuthorAccountIsInactive, ); pub const NSFW_USER_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DropNsfwUserAuthorRule", - |context| context.author_is_nsfw_user(), + |context| context.author().is_nsfw_user(), FilteredReason::ContainNsfwMedia, ); pub const NSFW_ADMIN_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( "DropNsfwAdminAuthorRule", - |context| context.author_is_nsfw_admin(), + |context| context.author().is_nsfw_admin(), FilteredReason::ContainNsfwMedia, ); @@ -71,9 +71,9 @@ impl Rule for ProtectedAuthorDropRule { } fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.author_is_protected() - && !context.is_author_viewer() - && (context.viewer_is_logged_out() || !context.viewer_follows_author()) + if context.author().is_protected() + && !context.viewer().is_author() + && (context.viewer().is_logged_out() || !context.viewer().follows_author()) { return VfAction::Drop(FilteredReason::AuthorIsProtected); } diff --git a/visibility-filtering/server.rs b/visibility-filtering/server.rs index 64f760e3..1ef69856 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,8 +28,11 @@ 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 async fn new( + datacenter: &str, + feature_switches: Arc, + ) -> Self { + crate::server_deps::build_prod_server(datacenter, feature_switches).await } pub fn from_endpoints( diff --git a/visibility-filtering/server_deps.rs b/visibility-filtering/server_deps.rs index 99b4fe3b..d59081d4 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -85,7 +85,10 @@ where } } -pub async fn build_prod_server(datacenter: &str) -> VFServer { +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; @@ -207,7 +210,11 @@ pub async fn build_prod_server(datacenter: &str) -> VFServer { safety_label_source.clone(), fallback_cache_mode, ); - let policies = crate::rules::Policies::new(); + let gating_countries = Arc::new(crate::params::NsfwGatingCountries::new()); + 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 policies = crate::rules::Policies::with_nsfw_gating_countries(gating_countries); let (home_rule_count, recommendations_rule_count) = policies.rule_counts(); let filter_tweets = FilterTweets::new(hydration_pipeline, policies); diff --git a/visibility-filtering/twemcache/host_pool.rs b/visibility-filtering/twemcache/host_pool.rs index d21ab4f5..555f42fa 100644 --- a/visibility-filtering/twemcache/host_pool.rs +++ b/visibility-filtering/twemcache/host_pool.rs @@ -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, } @@ -187,7 +187,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,7 +216,7 @@ 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; @@ -236,7 +236,7 @@ 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(); g.health.is_live() && !g.has_living_conn() @@ -279,14 +279,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(); 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; } @@ -331,7 +331,7 @@ 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(); if let Some(c) = &g.conn { 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, From 7ba776848b12d8422eb0f291ee03ea5c17ab0188 Mon Sep 17 00:00:00 2001 From: CI agent Date: Tue, 1 Sep 2026 20:47:22 +0000 Subject: [PATCH 13/18] Open-source X Recommendation Algorithm --- grox/core/data_loaders/kafka_loader.py | 16 - grox/flows/mm_emb/task_embedding_pub.py | 55 +- home-mixer/params/param.rs | 4 +- phoenix/crates/common/xai-recsys/src/util.rs | 82 +- .../serving/xai-recsys-engine/src/python.rs | 12 +- .../src/mm_embedding_client.rs | 41 +- .../xai-recsys-proto/proto/recsys.proto | 2 + phoenix/pyproject.toml | 1 + .../common/xai-proto/proto/recsys.proto | 2 + .../xai_checkpointing/orbax_encrypted.py | 288 +++++ phoenix/xrex/configs/xrecsys.py | 2 + phoenix/xrex/driver/hooks.py | 9 +- phoenix/xrex/models/recsys_model.py | 37 +- phoenix/xrex/models/recsys_sid.py | 72 ++ phoenix/xrex/train/checkpoint_write.py | 6 + phoenix/xrex/train/misc.py | 5 + phoenix/xrex/train/recsys_bundle_export.py | 634 ++++++++++ phoenix/xrex/train/trainer.py | 11 +- phoenix/xrex/train/trainer_recsys.py | 53 + visibility-filtering/rules/author_rules.rs | 354 ++++++ visibility-filtering/rules/context.rs | 9 + visibility-filtering/rules/mod.rs | 23 +- visibility-filtering/rules/nsfw_age_gating.rs | 639 ---------- .../rules/nsfw_interstitial.rs | 242 ---- visibility-filtering/rules/nullcast_rule.rs | 79 -- visibility-filtering/rules/registry.rs | 223 ++-- visibility-filtering/rules/rule_spec.rs | 89 ++ .../rules/socialgraph_rules.rs | 301 ----- visibility-filtering/rules/tes_rules.rs | 446 ------- .../rules/tweet_flag_rules.rs | 114 -- .../rules/tweet_label_drops.rs | 327 ----- visibility-filtering/rules/tweet_rules.rs | 1125 +++++++++++++++++ .../rules/user_label_drops.rs | 239 ---- visibility-filtering/rules/user_rules.rs | 283 ----- 34 files changed, 2964 insertions(+), 2861 deletions(-) create mode 100644 phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py create mode 100644 phoenix/xrex/models/recsys_sid.py create mode 100644 phoenix/xrex/train/recsys_bundle_export.py create mode 100644 visibility-filtering/rules/author_rules.rs delete mode 100644 visibility-filtering/rules/nsfw_age_gating.rs delete mode 100644 visibility-filtering/rules/nsfw_interstitial.rs delete mode 100644 visibility-filtering/rules/nullcast_rule.rs create mode 100644 visibility-filtering/rules/rule_spec.rs delete mode 100644 visibility-filtering/rules/socialgraph_rules.rs delete mode 100644 visibility-filtering/rules/tes_rules.rs delete mode 100644 visibility-filtering/rules/tweet_flag_rules.rs delete mode 100644 visibility-filtering/rules/tweet_label_drops.rs create mode 100644 visibility-filtering/rules/tweet_rules.rs delete mode 100644 visibility-filtering/rules/user_label_drops.rs delete mode 100644 visibility-filtering/rules/user_rules.rs diff --git a/grox/core/data_loaders/kafka_loader.py b/grox/core/data_loaders/kafka_loader.py index 7cd2e47c..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 @@ -53,27 +52,12 @@ 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 self._max_qps_per_partition = self.loader_config.max_qps_per_partition self._limit_item: RateLimitItemPerSecond | 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 - def _is_shutdown(self) -> bool: try: return self._shutdown_event.is_set() 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/home-mixer/params/param.rs b/home-mixer/params/param.rs index 713dda5f..dc563640 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-31T16:17:17Z +// mirrored from config feature-switch defaults; last sync 2026-09-01T16:42:25Z use xai_feature_switches::param; param!( @@ -909,7 +909,7 @@ param!( AdsBlenderType, String, "rust_home_mixer_ads_blender", - "partition_organic_low_risk" + "multi_risk" ); param!( diff --git a/phoenix/crates/common/xai-recsys/src/util.rs b/phoenix/crates/common/xai-recsys/src/util.rs index 99a95c90..ceb672e7 100644 --- a/phoenix/crates/common/xai-recsys/src/util.rs +++ b/phoenix/crates/common/xai-recsys/src/util.rs @@ -331,6 +331,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, @@ -424,23 +448,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); - } - } else { + 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 {provided_dim} != model {search_query_embedding_dim}; leaving zeros" + "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; @@ -1732,6 +1752,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 diff --git a/phoenix/crates/serving/xai-recsys-engine/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index 74bbfda7..8cd7bb47 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -2046,8 +2046,16 @@ impl PrepareBatch for RankingBatchPrep { .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, + ); + } } }); } 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 11082080..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 @@ -216,23 +216,28 @@ impl MmEmbeddingsClient { } } - pub fn fetch_mm_embeddings_sync( + 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 processed = total_posts.min(dest.len() / emb_dim); let missing = AtomicUsize::new(0); - results + dest[..processed * emb_dim] .par_chunks_mut(emb_dim) .zip(post_ids.par_iter().take(processed)) .for_each(|(dst, &post_id)| { @@ -255,6 +260,18 @@ 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) } @@ -417,6 +434,14 @@ mod tests { 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] diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index 17e6a5c2..e08447ef 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -453,6 +453,7 @@ enum ActionName { 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; @@ -1359,6 +1360,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 { 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-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index 17e6a5c2..e08447ef 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -453,6 +453,7 @@ enum ActionName { 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; @@ -1359,6 +1360,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 { 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..aa0aee45 --- /dev/null +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py @@ -0,0 +1,288 @@ +# 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 _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 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, handler_kwargs): + 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 + + 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) + + gb = handler_kwargs.get("save_concurrent_gb") + if gb is not None: + return EncryptedArrayHandler(int(gb) * 10**9) + return EncryptedArrayHandler() + + +def encrypted_checkpointer( + kms_client, encryption_chunk_size: int, timeout_secs: int, array_handler_cls, 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, handler_kwargs) + 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=handler_kwargs.get("use_zarr3", False), + save_concurrent_bytes=_pytree._concurrent_bytes(handler_kwargs.get("save_concurrent_gb")), + restore_concurrent_bytes=_pytree._concurrent_bytes( + 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, **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 + + +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 + if _ENCRYPTED_CHECKPOINTER is None: + handler_kwargs: dict = {"use_zarr3": True} + if save_concurrent_gb is not None: + handler_kwargs["save_concurrent_gb"] = save_concurrent_gb + handler_kwargs["restore_concurrent_gb"] = save_concurrent_gb + _ENCRYPTED_SAVE_CONCURRENT_GB = save_concurrent_gb + _ENCRYPTED_TIMEOUT_SECS = timeout_secs + _ENCRYPTED_CHECKPOINTER = encrypted_checkpointer( + kms_client, encryption_chunk_size, timeout_secs, array_handler_cls, 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, + ) + return _ENCRYPTED_CHECKPOINTER + + +def assert_all_enveloped(path: pathlib.Path) -> None: + 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/xrex/configs/xrecsys.py b/phoenix/xrex/configs/xrecsys.py index 34220e3a..cbaac72e 100644 --- a/phoenix/xrex/configs/xrecsys.py +++ b/phoenix/xrex/configs/xrecsys.py @@ -628,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, diff --git a/phoenix/xrex/driver/hooks.py b/phoenix/xrex/driver/hooks.py index fc8cee6a..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() diff --git a/phoenix/xrex/models/recsys_model.py b/phoenix/xrex/models/recsys_model.py index 3944837d..e23c2053 100644 --- a/phoenix/xrex/models/recsys_model.py +++ b/phoenix/xrex/models/recsys_model.py @@ -59,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, @@ -542,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 @@ -614,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 @@ -2642,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, @@ -2660,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, @@ -2672,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" @@ -2704,8 +2735,10 @@ 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 ) 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/train/checkpoint_write.py b/phoenix/xrex/train/checkpoint_write.py index 6260ae50..f5b0f3e8 100644 --- a/phoenix/xrex/train/checkpoint_write.py +++ b/phoenix/xrex/train/checkpoint_write.py @@ -66,6 +66,12 @@ def save_checkpoint( path = self.get_checkpoint_path(self.ctx) + if getattr(self, "_restored_encrypted", False): + raise ValueError( + "restored from an ENCRYPTED checkpoint but the orbax writer cannot " + "produce encrypted saves — refusing to downgrade to plaintext saves" + ) + with tracer.start_as_current_span("write_checksum"): if self.checkpoint_config.write_checksums: if checksum_dict is None: diff --git a/phoenix/xrex/train/misc.py b/phoenix/xrex/train/misc.py index d0b08db3..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 @@ -52,6 +53,10 @@ 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 diff --git a/phoenix/xrex/train/recsys_bundle_export.py b/phoenix/xrex/train/recsys_bundle_export.py new file mode 100644 index 00000000..2ceb1070 --- /dev/null +++ b/phoenix/xrex/train/recsys_bundle_export.py @@ -0,0 +1,634 @@ +# 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" + + +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.float32) + cont_preds = candidate_continuous_predictions.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 512a6b28..5c74d583 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -1166,6 +1166,9 @@ def warm_start_staging_spec(self): 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.encrypt or self.checkpoint_config.save_method == "tensorstore" + def maybe_load_checkpoint( self, ctx: TrainerContext, tag: str | None = None ) -> tuple[bool, int, int]: @@ -1173,6 +1176,12 @@ 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() @@ -1180,7 +1189,7 @@ def maybe_load_checkpoint( use_streamed_restore = ( self.checkpoint_config.restore_streamed and ctx.checkpoint.format == "orbax" - and self.checkpoint_config.save_method != "tensorstore" + and not self._uses_tensorstore_save() ) if self.checkpoint_config.restore_streamed and not use_streamed_restore: rank_logger.info( diff --git a/phoenix/xrex/train/trainer_recsys.py b/phoenix/xrex/train/trainer_recsys.py index 2c53a37f..fe357236 100644 --- a/phoenix/xrex/train/trainer_recsys.py +++ b/phoenix/xrex/train/trainer_recsys.py @@ -415,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 @@ -453,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 @@ -2824,6 +2832,46 @@ 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: + try: + for bundle_file in bundle_files or (): + 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(bundle_file.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]], @@ -2831,6 +2879,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() @@ -2899,6 +2948,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) @@ -3173,6 +3223,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, @@ -3180,6 +3232,7 @@ def save_checkpoint(self, *args, **kwargs): data_pos, prefix, stub, + bundle_files, ) if store != Store.FS: diff --git a/visibility-filtering/rules/author_rules.rs b/visibility-filtering/rules/author_rules.rs new file mode 100644 index 00000000..758ca0bd --- /dev/null +++ b/visibility-filtering/rules/author_rules.rs @@ -0,0 +1,354 @@ +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, HydratedTweetCandidate, VfAction, ViewerAuthorRelationship, ViewerFeatures, + }; + use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; + use crate::rules::{test_context, Rule}; + + 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() + ); + } + + 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() + ); + } + + 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 index b2c60141..3900bae1 100644 --- a/visibility-filtering/rules/context.rs +++ b/visibility-filtering/rules/context.rs @@ -1,4 +1,5 @@ use crate::models::{HydratedTweetCandidate, SafetyLabelType, ViewerFeatures}; +use crate::params::NsfwGatingCountries; use crate::rules::registry::SafetyLevel; use xai_core_entities::entities::TakedownReason; use xai_x_thrift::user_labels::LabelValue; @@ -7,6 +8,7 @@ pub struct RuleContext<'a> { safety_level: SafetyLevel, viewer: &'a ViewerFeatures, candidate: &'a HydratedTweetCandidate, + nsfw_gating_countries: &'a NsfwGatingCountries, } impl<'a> RuleContext<'a> { @@ -14,11 +16,13 @@ impl<'a> RuleContext<'a> { safety_level: SafetyLevel, viewer: &'a ViewerFeatures, candidate: &'a HydratedTweetCandidate, + nsfw_gating_countries: &'a NsfwGatingCountries, ) -> Self { Self { safety_level, viewer, candidate, + nsfw_gating_countries, } } @@ -46,6 +50,11 @@ impl<'a> RuleContext<'a> { 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)] diff --git a/visibility-filtering/rules/mod.rs b/visibility-filtering/rules/mod.rs index 33f7df2a..fa5ab79a 100644 --- a/visibility-filtering/rules/mod.rs +++ b/visibility-filtering/rules/mod.rs @@ -1,19 +1,13 @@ +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::VfAction; use xai_visibility_filtering::models::FilteredReason; @@ -72,7 +66,16 @@ pub(crate) fn test_context<'a>( viewer: &'a crate::models::ViewerFeatures, candidate: &'a crate::models::HydratedTweetCandidate, ) -> RuleContext<'a> { - RuleContext::new(SafetyLevel::TimelineHome, viewer, candidate) + use std::sync::LazyLock; + + static NSFW_GATING_COUNTRIES: LazyLock = + LazyLock::new(crate::params::NsfwGatingCountries::new); + RuleContext::new( + SafetyLevel::TimelineHome, + viewer, + candidate, + &NSFW_GATING_COUNTRIES, + ) } #[cfg(test)] diff --git a/visibility-filtering/rules/nsfw_age_gating.rs b/visibility-filtering/rules/nsfw_age_gating.rs deleted file mode 100644 index 3c4ab4cf..00000000 --- a/visibility-filtering/rules/nsfw_age_gating.rs +++ /dev/null @@ -1,639 +0,0 @@ -use crate::models::{SafetyLabelType, VfAction}; -use crate::params::NsfwGatingCountries; -use crate::rules::{Rule, RuleContext}; -use std::sync::Arc; -use xai_visibility_filtering::models::FilteredReason; - -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) -} - -pub struct SensitiveViewerLoggedOutDropRule; - -impl Rule for SensitiveViewerLoggedOutDropRule { - fn name(&self) -> &'static str { - "SensitiveViewerLoggedOutDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.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 { - if context.viewer().is_underage() && sensitive_base_condition(context) { - VfAction::Drop(FilteredReason::ContainNsfwMedia) - } else { - VfAction::Allow - } - } -} - -pub struct SensitiveViewerNoStatedAgeDropRule { - gating_countries: Arc, -} - -impl SensitiveViewerNoStatedAgeDropRule { - pub fn new(gating_countries: Arc) -> Self { - Self { gating_countries } - } -} - -impl Rule for SensitiveViewerNoStatedAgeDropRule { - fn name(&self) -> &'static str { - "SensitiveViewerNoStatedAgeDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer().has_no_stated_age() - && context - .viewer() - .country() - .is_some_and(|c| self.gating_countries.contains(c)) - && sensitive_base_condition(context) - { - VfAction::Drop(FilteredReason::ContainNsfwMedia) - } else { - VfAction::Allow - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, NsfwFeature, Viewer, ViewerAge, ViewerFeatures, - }; - use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; - - fn gating_viewer(age: ViewerAge) -> ViewerFeatures { - ViewerFeatures { - viewer_age: age, - country_code: Some("de".into()), - ..viewer(VIEWER_ID) - } - } - - fn media_candidate_with_label(label: SafetyLabelType) -> HydratedTweetCandidate { - candidate().with_label(label).with_media().build() - } - - fn no_stated_age_rule() -> SensitiveViewerNoStatedAgeDropRule { - SensitiveViewerNoStatedAgeDropRule::new(Arc::new(NsfwGatingCountries::new())) - } - - fn nsfw_author_media_candidate() -> HydratedTweetCandidate { - candidate() - .with_media() - .with_author_features(AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }) - .build() - } - - #[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( - &gating_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( - &gating_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( - &gating_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, - ..gating_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( - &gating_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( - &gating_viewer(ViewerAge::Unknown), - &c - )), - VfAction::Allow - )); - assert!(matches!( - no_stated_age_rule().evaluate(&crate::rules::test_context( - &gating_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( - &gating_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( - &gating_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!( - no_stated_age_rule().evaluate(&crate::rules::test_context( - &gating_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()), - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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, - ..gating_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, - ..gating_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( - &gating_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( - &gating_viewer(ViewerAge::Unknown), - &c - )), - VfAction::Allow - )); - assert!(matches!( - no_stated_age_rule().evaluate(&crate::rules::test_context( - &gating_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 = VIEWER_ID; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &gating_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( - &gating_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 = VIEWER_ID; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &gating_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!( - no_stated_age_rule().evaluate(&crate::rules::test_context( - &gating_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()), - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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, - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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()), - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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()), - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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, - ..gating_viewer(ViewerAge::NotStated) - }; - assert!(matches!( - no_stated_age_rule().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( - &gating_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( - &gating_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( - &gating_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( - &gating_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( - &gating_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( - &gating_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( - &gating_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 = VIEWER_ID; - assert!(matches!( - SensitiveViewerUnderageDropRule.evaluate(&crate::rules::test_context( - &gating_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( - &gating_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( - &gating_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, - ..gating_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, - ..gating_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, - ..gating_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( - &gating_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 6438c02f..00000000 --- a/visibility-filtering/rules/nsfw_interstitial.rs +++ /dev/null @@ -1,242 +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.tweet().has_safety_label(self.label) - && !context.viewer().is_author() - && !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 { - if context.tweet().is_nsfw_flagged() - && context.tweet().has_media() - && !context.viewer().is_author() - && !context.viewer().allows_sensitive_media() - { - return VfAction::Interstitial(FilteredReason::ContainNsfwMedia); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{AuthorFeatures, HydratedTweetCandidate, NsfwFeature}; - use crate::rules::fixtures::{ - author_viewer, candidate, sensitive_opt_in_viewer, viewer, VIEWER_ID, - }; - - #[test] - fn interstitial_blurs_non_opt_in() { - let c = candidate() - .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) - .build(); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Interstitial(_) - )); - } - - #[test] - fn interstitial_allows_opt_in() { - let c = candidate() - .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) - .build(); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&sensitive_opt_in_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn interstitial_allows_self_view() { - let c = candidate() - .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) - .build(); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } - - #[test] - fn interstitial_allows_no_label() { - let c = candidate().build(); - assert!(matches!( - NSFW_HIGH_PRECISION_INTERSTITIAL - .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - fn nsfw_author_candidate() -> HydratedTweetCandidate { - candidate() - .with_media() - .with_author_features(AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }) - .build() - } - - #[test] - fn author_interstitial_blurs_non_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &viewer(VIEWER_ID), - &nsfw_author_candidate() - )), - VfAction::Interstitial(_) - )); - } - - #[test] - fn author_interstitial_allows_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &sensitive_opt_in_viewer(), - &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(VIEWER_ID), &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(VIEWER_ID), - &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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn tweet_flag_interstitial_allows_opt_in() { - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context( - &sensitive_opt_in_viewer(), - &nsfw_tweet_flag_candidate() - )), - VfAction::Allow - )); - } - - #[test] - fn tweet_flag_interstitial_allows_self_view() { - let c = nsfw_tweet_flag_candidate(); - assert!(matches!( - NsfwAuthorInterstitialRule.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/nullcast_rule.rs b/visibility-filtering/rules/nullcast_rule.rs deleted file mode 100644 index c677ab66..00000000 --- a/visibility-filtering/rules/nullcast_rule.rs +++ /dev/null @@ -1,79 +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 { - if context.tweet().is_nullcast() - && !context.tweet().is_retweet() - && !context.tweet().is_community_tweet() - { - return VfAction::Drop(FilteredReason::TweetIsNullcast); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{HydratedTweetCandidate, TweetFeatures}; - use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; - - fn nullcast_candidate() -> HydratedTweetCandidate { - candidate() - .with_tweet_features(TweetFeatures { - is_nullcast: true, - ..Default::default() - }) - .build() - } - - #[test] - fn nullcast_non_retweet_drops() { - let rule = NullcastedTweetDropRule; - let c = nullcast_candidate(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn nullcast_community_tweet_allows() { - let rule = NullcastedTweetDropRule; - let mut c = nullcast_candidate(); - c.tweet_features.is_community_tweet = true; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn nullcast_retweet_allows() { - let rule = NullcastedTweetDropRule; - let mut c = nullcast_candidate(); - c.tweet_features.core.source_tweet_id = Some(99); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn non_nullcast_allows() { - let rule = NullcastedTweetDropRule; - let c = candidate().build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index 5c7d204b..fdeb40b9 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -1,28 +1,9 @@ -use crate::models::{HydratedTweetCandidate, VfAction, ViewerFeatures}; +use crate::models::{HydratedTweetCandidate, ViewerFeatures}; use crate::params::NsfwGatingCountries; -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::rule_spec::RuleSpec; +use crate::rules::{author_rules, tweet_rules}; use crate::rules::{evaluate_rules, Rule, RuleContext, Verdict}; use std::sync::Arc; -use xai_visibility_filtering::models::FilteredReason; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum SafetyLevel { @@ -45,6 +26,7 @@ pub struct Policies { filter_all: Vec>, timeline_home: Vec>, timeline_home_recommendations: Vec>, + nsfw_gating_countries: Arc, } impl Policies { @@ -54,9 +36,10 @@ impl Policies { pub fn with_nsfw_gating_countries(gating_countries: Arc) -> Self { Self { - filter_all: vec![Box::new(FilterAllRule)], - timeline_home: timeline_home_policy(&gating_countries), - timeline_home_recommendations: timeline_home_recommendations_policy(&gating_countries), + filter_all: rule_specs(tweet_rules::FILTER_ALL).collect(), + timeline_home: timeline_home_policy(), + timeline_home_recommendations: timeline_home_recommendations_policy(), + nsfw_gating_countries: gating_countries, } } @@ -74,7 +57,7 @@ impl Policies { viewer: &ViewerFeatures, candidate: &HydratedTweetCandidate, ) -> Verdict { - let context = RuleContext::new(level, viewer, candidate); + let context = RuleContext::new(level, viewer, candidate, &self.nsfw_gating_countries); evaluate_rules(self.select(level), &context) } @@ -97,98 +80,48 @@ impl Default for Policies { } } -struct FilterAllRule; - -impl Rule for FilterAllRule { - fn name(&self) -> &'static str { - "FilterAllRule" - } - - fn evaluate(&self, _context: &RuleContext<'_>) -> VfAction { - VfAction::Drop(FilteredReason::UnspecifiedReason) - } +fn rule_specs(specs: &'static [RuleSpec]) -> impl Iterator> { + specs + .iter() + .map(|spec| Box::new(spec.clone()) as Box) } -fn base_home_rules(gating_countries: &Arc) -> 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::new(Arc::clone( - gating_countries, - ))), - 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 base_home_rules() -> Vec> { + let mut rules: Vec> = Vec::new(); + rules.extend(rule_specs(author_rules::AUTHOR_STATE_DROPS)); + rules.extend(rule_specs(author_rules::SOCIALGRAPH_DROPS)); + rules.extend(rule_specs(tweet_rules::TWEET_LABEL_DROPS)); + rules.extend(rule_specs(tweet_rules::NULLCAST_DROP)); + rules.extend(rule_specs(tweet_rules::TES_HOME_DROPS)); + rules.extend(rule_specs(tweet_rules::SENSITIVE_VIEWER_DROPS)); + rules.extend(rule_specs(tweet_rules::EXCLUSIVE_TWEET_DROP)); + rules.extend(rule_specs(tweet_rules::NSFW_MEDIA_INTERSTITIALS)); + rules.extend(rule_specs(tweet_rules::NSFW_AUTHOR_INTERSTITIAL)); + rules } -fn timeline_home_policy(gating_countries: &Arc) -> Vec> { - base_home_rules(gating_countries) +fn timeline_home_policy() -> Vec> { + base_home_rules() } -fn timeline_home_recommendations_policy( - gating_countries: &Arc, -) -> Vec> { - let mut rules = base_home_rules(gating_countries); - 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); +fn timeline_home_recommendations_policy() -> Vec> { + let mut rules = base_home_rules(); + rules.extend(rule_specs(tweet_rules::RECS_MEDIA_DROPS)); + rules.extend(rule_specs(author_rules::OON_NSFW_AUTHOR_DROPS)); + rules.extend(rule_specs(tweet_rules::OON_TWEET_FLAG_DROPS)); + rules.extend(rule_specs(tweet_rules::OON_TWEET_LABEL_DROPS)); + rules.extend(rule_specs(author_rules::OON_USER_LABEL_DROPS)); rules } #[cfg(test)] mod tests { use super::*; - use crate::models::{HydratedTweetCandidate, MediaFeature, TweetFeatures, ViewerFeatures}; + use crate::models::{ + HydratedTweetCandidate, MediaFeature, TweetFeatures, VfAction, ViewerFeatures, + }; use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; + use xai_visibility_filtering::models::FilteredReason; struct RecommendationsOnlyRule; @@ -213,6 +146,7 @@ mod tests { filter_all: vec![Box::new(RecommendationsOnlyRule)], timeline_home: vec![Box::new(RecommendationsOnlyRule)], timeline_home_recommendations: vec![Box::new(RecommendationsOnlyRule)], + nsfw_gating_countries: Arc::new(NsfwGatingCountries::new()), }; let viewer = ViewerFeatures::default(); let candidate = HydratedTweetCandidate::default(); @@ -277,12 +211,89 @@ rust_vf: fn filter_all_rule_drops_even_self_view() { let candidate = candidate().build(); let viewer = author_viewer(); + let spec = &tweet_rules::FILTER_ALL[0]; assert!(matches!( - FilterAllRule.evaluate(&crate::rules::test_context(&viewer, &candidate)), + spec.evaluate(&crate::rules::test_context(&viewer, &candidate)), VfAction::Drop(_) )); } + #[test] + fn wired_rule_order_matches_pre_migration_sequence() { + let policies = Policies::new(); + assert_eq!( + policies.wired_rule_names(SafetyLevel::FilterAll), + vec!["FilterAllRule"] + ); + let home = policies.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!( + policies.wired_rule_names(SafetyLevel::TimelineHomeRecommendations), + recs + ); + } + #[test] fn filter_all_policy_drops_pristine_candidate() { let policies = Policies::new(); diff --git a/visibility-filtering/rules/rule_spec.rs b/visibility-filtering/rules/rule_spec.rs new file mode 100644 index 00000000..b43c1987 --- /dev/null +++ b/visibility-filtering/rules/rule_spec.rs @@ -0,0 +1,89 @@ +use crate::models::VfAction; +use crate::rules::context::{AuthorPredicates, TweetPredicates}; +use crate::rules::{Rule, RuleContext}; +use xai_visibility_filtering::models::FilteredReason; + +#[derive(Clone)] +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, + }, +} + +#[derive(Clone)] +pub(super) enum RuleAction { + Drop(FilteredReason), + SensitiveMediaInterstitial(FilteredReason), +} + +impl Rule for RuleSpec { + fn name(&self) -> &'static str { + match self { + RuleSpec::Tweet { name, .. } + | RuleSpec::Author { name, .. } + | RuleSpec::Custom { name, .. } => name, + } + } + + 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 79815221..00000000 --- a/visibility-filtering/rules/socialgraph_rules.rs +++ /dev/null @@ -1,301 +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 { - if context.viewer().is_logged_out() { - return VfAction::Allow; - } - if context.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 { - 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 struct ViewerMutesAuthorRule; - -impl Rule for ViewerMutesAuthorRule { - fn name(&self) -> &'static str { - "ViewerMutesAuthorRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.viewer().is_logged_out() { - return VfAction::Allow; - } - if context.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 { - 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) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - ExclusiveContentFeatures, HydratedTweetCandidate, ViewerAuthorRelationship, - }; - use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; - - 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 - } - - fn candidate_with_relationship( - relationship: ViewerAuthorRelationship, - ) -> HydratedTweetCandidate { - candidate().with_relationship(relationship).build() - } - - #[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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &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 c = candidate() - .with_relationship(ViewerAuthorRelationship { - viewer_mutes_retweets_from_author: true, - ..Default::default() - }) - .retweet_of(99) - .build(); - assert!(matches!( - MutedRetweetsRule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn muted_retweets_allows_logged_out_viewer() { - let c = candidate() - .with_relationship(ViewerAuthorRelationship { - viewer_mutes_retweets_from_author: true, - ..Default::default() - }) - .retweet_of(99) - .build(); - 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(VIEWER_ID), - &candidate().build(), - )); - 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(&author_viewer(), &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 35df980f..00000000 --- a/visibility-filtering/rules/tes_rules.rs +++ /dev/null @@ -1,446 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -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 { - if context.tweet().is_stale() && !context.tweet().is_retweet() { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -pub struct DropLegalTakendownPostRule; - -impl Rule for DropLegalTakendownPostRule { - fn name(&self) -> &'static str { - "DropLegalTakendownPostRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if !context.viewer().is_author() && context.takedown().legal_in_viewer_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 !context.viewer().is_author() && context.takedown().local_laws_in_viewer_country() { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -pub struct DropTweetsWithGeoRestrictedMediaRule; - -impl Rule for DropTweetsWithGeoRestrictedMediaRule { - fn name(&self) -> &'static str { - "DropTweetsWithGeoRestrictedMediaRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.takedown().media_restricted_in_viewer_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 { - if context.tweet().has_dmca_media() { - return VfAction::Drop(FilteredReason::UnspecifiedReason); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{ - HydratedTweetCandidate, MediaFeature, TakedownFeature, TweetFeatures, ViewerFeatures, - }; - use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; - use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; - - fn takedown_candidate(reasons: Vec) -> HydratedTweetCandidate { - candidate() - .with_tweet_features(TweetFeatures { - takedown: TakedownFeature { - reasons, - ..Default::default() - }, - ..Default::default() - }) - .build() - } - - 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 = candidate() - .with_tweet_features(TweetFeatures { - edit_control: stale_edit_control(), - ..Default::default() - }) - .build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn non_stale_tweet_allows() { - let rule = DropStaleTweetsRule; - let c = candidate().build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn stale_retweet_allows() { - let rule = DropStaleTweetsRule; - let c = candidate() - .with_tweet_features(TweetFeatures { - edit_control: stale_edit_control(), - ..Default::default() - }) - .retweet_of(99) - .build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - fn viewer_with_country(country: &str) -> ViewerFeatures { - ViewerFeatures { - country_code: Some(country.to_string()), - ..viewer(VIEWER_ID) - } - } - - #[test] - fn takedown_drops_in_matching_country() { - let rule = DropLegalTakendownPostRule; - let c = takedown_candidate(vec![ - TakedownReason::LegalRequest { - country_code: "de".to_string(), - }, - TakedownReason::UnspecifiedReason { - country_code: "fr".to_string(), - }, - ]); - 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 = takedown_candidate(vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }]); - 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 = takedown_candidate(vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }]); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn legal_rule_ignores_local_laws_countries() { - let rule = DropLegalTakendownPostRule; - let c = takedown_candidate(vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }]); - 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 = takedown_candidate(vec![ - TakedownReason::BystanderReport { - country_code: "de".to_string(), - }, - TakedownReason::BystanderReport { - country_code: "fr".to_string(), - }, - ]); - 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 = takedown_candidate(vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }]); - 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 = takedown_candidate(vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }]); - 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 mut c = takedown_candidate(vec![TakedownReason::LegalRequest { - country_code: "de".to_string(), - }]); - c.author_id = VIEWER_ID; - 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 mut c = takedown_candidate(vec![TakedownReason::BystanderReport { - country_code: "de".to_string(), - }]); - c.author_id = VIEWER_ID; - 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 = takedown_candidate(vec![ - TakedownReason::Dmca, - TakedownReason::HatefulImagery, - TakedownReason::Unknown, - ]); - 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 { - 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 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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn geo_restricted_drops_even_for_author() { - let rule = DropTweetsWithGeoRestrictedMediaRule; - let mut c = geo_candidate(&[], &["de"]); - c.author_id = VIEWER_ID; - 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 = candidate() - .with_tweet_features(TweetFeatures { - media: MediaFeature { - has_dmca_media: true, - ..Default::default() - }, - ..Default::default() - }) - .build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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 06d6d233..00000000 --- a/visibility-filtering/rules/tweet_flag_rules.rs +++ /dev/null @@ -1,114 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -#[derive(Clone)] -pub struct TweetFlagDropRule { - name: &'static str, - flag: fn(&RuleContext<'_>) -> bool, - reason: FilteredReason, -} - -impl TweetFlagDropRule { - pub const fn new( - name: &'static str, - flag: fn(&RuleContext<'_>) -> 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) { - return VfAction::Drop(self.reason.clone()); - } - VfAction::Allow - } -} - -pub const TWEET_NSFW_USER_DROP: TweetFlagDropRule = TweetFlagDropRule::new( - "TweetNsfwUserDropRule", - |context| context.tweet().has_nsfw_user_flag(), - FilteredReason::ContainNsfwMedia, -); -pub const TWEET_NSFW_ADMIN_DROP: TweetFlagDropRule = TweetFlagDropRule::new( - "TweetNsfwAdminDropRule", - |context| context.tweet().has_nsfw_admin_flag(), - FilteredReason::ContainNsfwMedia, -); - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{HydratedTweetCandidate, NsfwFeature, TweetFeatures}; - use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; - - fn candidate_with_nsfw_flags(user: bool, admin: bool) -> HydratedTweetCandidate { - candidate() - .with_tweet_features(TweetFeatures { - nsfw: NsfwFeature { user, admin }, - ..Default::default() - }) - .build() - } - - #[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(VIEWER_ID), &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(VIEWER_ID), &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(&author_viewer(), &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(VIEWER_ID), &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(VIEWER_ID), &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(&author_viewer(), &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 32bcab93..00000000 --- a/visibility-filtering/rules/tweet_label_drops.rs +++ /dev/null @@ -1,327 +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.tweet().has_safety_label(self.label) { - return VfAction::Allow; - } - if self.exempt_author && context.viewer().is_author() { - 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::rules::fixtures::{ - author_viewer, candidate, sensitive_opt_in_viewer, viewer, VIEWER_ID, - }; - - #[test] - fn drops_non_author_with_mapped_reason() { - let c = candidate().with_label(SafetyLabelType::PDNA).build(); - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::SafetyResult(_)) - )); - - let c = candidate().with_label(SafetyLabelType::BOUNCE).build(); - assert!(matches!( - BOUNCE_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::TweetIsBounced) - )); - - let c = candidate().with_label(SafetyLabelType::SPAM).build(); - assert!(matches!( - SPAM_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - - let c = candidate() - .with_label(SafetyLabelType::NSFW_HIGH_RECALL) - .build(); - assert!(matches!( - NSFW_HIGH_RECALL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn author_exempt_rules_allow_self_view() { - let c = candidate().with_label(SafetyLabelType::PDNA).build(); - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - let c = candidate() - .with_label(SafetyLabelType::DO_NOT_AMPLIFY) - .build(); - 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) - .build(); - 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 c = candidate() - .with_label(SafetyLabelType::NSFW_HIGH_PRECISION) - .build(); - assert!(matches!( - NSFW_HIGH_PRECISION_DROP - .evaluate(&crate::rules::test_context(&sensitive_opt_in_viewer(), &c)), - VfAction::Drop(FilteredReason::ContainNsfwMedia) - )); - } - - #[test] - fn no_label_allows() { - let c = candidate().build(); - assert!(matches!( - PDNA_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[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).build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate().with_label(rule.label).followed().build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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).build(); - 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) - .build(); - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&author_viewer(), &c)), - VfAction::Allow - )); - let c = candidate().build(); - assert!(matches!( - MALICIOUS_URL_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn fosnr_abuse_insults_oon_drops_all_non_authors() { - let c = candidate() - .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) - .build(); - assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP - .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate() - .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) - .followed() - .build(); - assert!(matches!( - FOSNR_ABUSE_INSULTS_OON_DROP - .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::PossiblyUndesirable) - )); - let c = candidate() - .with_label(SafetyLabelType::FOSNR_ABUSE_INSULTS) - .build(); - 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..7a104be3 --- /dev/null +++ b/visibility-filtering/rules/tweet_rules.rs @@ -0,0 +1,1125 @@ +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() && 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, TakedownFeature, TweetFeatures, VfAction, Viewer, ViewerAge, ViewerFeatures, + }; + use crate::rules::fixtures::{ + author_viewer, candidate, logged_out_viewer, sensitive_opt_in_viewer, viewer, VIEWER_ID, + }; + use crate::rules::{test_context, Rule, RuleContext}; + use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; + + 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() + ); + } + + 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() + ); + } + + 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 custom_drop(_context: &RuleContext<'_>) -> VfAction { + VfAction::Drop(FilteredReason::UnspecifiedReason) + } + + fn custom_allow(_context: &RuleContext<'_>) -> VfAction { + VfAction::Allow + } + + fn custom_drop_even_self_view(context: &RuleContext<'_>) -> VfAction { + if context.tweet().is_nullcast() { + VfAction::Drop(FilteredReason::TweetIsNullcast) + } else { + VfAction::Allow + } + } + + #[test] + fn custom_row_returns_leaf_action() { + let drop_row = RuleSpec::Custom { + name: "CustomDropLeaf", + evaluate: custom_drop, + }; + let allow_row = RuleSpec::Custom { + name: "CustomAllowLeaf", + evaluate: custom_allow, + }; + let pristine = candidate().build(); + assert_drops( + &drop_row, + &viewer(VIEWER_ID), + &pristine, + &FilteredReason::UnspecifiedReason, + ); + assert_drops( + &drop_row, + &author_viewer(), + &pristine, + &FilteredReason::UnspecifiedReason, + ); + assert_allows(&allow_row, &viewer(VIEWER_ID), &pristine); + } + + 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 + } + + fn nsfw_flag_media_candidates() -> Vec { + let author_user = candidate() + .with_media() + .with_author_features(AuthorFeatures { + is_nsfw_user: true, + ..Default::default() + }) + .build(); + let tweet_user = { + let mut c = author_user.clone(); + c.author_features = AuthorFeatures::default(); + c.tweet_features.nsfw = NsfwFeature { + user: true, + admin: false, + }; + c + }; + let tweet_admin = { + let mut c = tweet_user.clone(); + c.tweet_features.nsfw = NsfwFeature { + user: false, + admin: true, + }; + c + }; + let both = { + let mut c = tweet_user.clone(); + c.author_features.is_nsfw_admin = true; + c + }; + vec![author_user, tweet_user, tweet_admin, both] + } + + #[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().remove(0); + no_media.tweet_features.media.has_media = false; + assert_allows(spec, &viewer(VIEWER_ID), &no_media); + let mut no_flags = nsfw_flag_media_candidates().remove(1); + 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); + + 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: TakedownFeature { + reasons, + ..Default::default() + }, + ..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::Dmca, + TakedownReason::HatefulImagery, + TakedownReason::Unknown, + ]); + assert_allows(legal, &viewer_with_country("de"), &non_country); + assert_allows(local, &viewer_with_country("de"), &non_country); + } + + #[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); + } + + #[test] + fn custom_row_does_not_add_author_exemption() { + let spec = RuleSpec::Custom { + name: "CustomNullcastLeaf", + evaluate: custom_drop_even_self_view, + }; + let firing = candidate() + .with_tweet_features(TweetFeatures { + is_nullcast: true, + ..Default::default() + }) + .build(); + assert_drops( + &spec, + &author_viewer(), + &firing, + &FilteredReason::TweetIsNullcast, + ); + assert_allows(&spec, &author_viewer(), &candidate().build()); + } +} diff --git a/visibility-filtering/rules/user_label_drops.rs b/visibility-filtering/rules/user_label_drops.rs deleted file mode 100644 index e72c82e7..00000000 --- a/visibility-filtering/rules/user_label_drops.rs +++ /dev/null @@ -1,239 +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.viewer().is_author() { - return VfAction::Allow; - } - if !context.author().has_user_label(self.label) { - return VfAction::Allow; - } - if self.require_non_follower - && !context.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::HydratedTweetCandidate; - use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; - - fn candidate_with_user_label(label: LabelValue) -> HydratedTweetCandidate { - candidate().with_author_user_label(label).build() - } - - #[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(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - let c = candidate_with_user_label(LabelValue::COMPROMISED); - assert!(matches!( - COMPROMISED_USER_DROP.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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(VIEWER_ID), &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 = candidate().build(); - assert!(matches!( - IMPERSONATION_HIGH_PRECISION_USER_DROP - .evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Drop(FilteredReason::UnspecifiedReason) - )); - - assert!(matches!( - ABUSIVE_HIGH_RECALL_USER_DROP - .evaluate(&crate::rules::test_context(&logged_out_viewer(), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } -} diff --git a/visibility-filtering/rules/user_rules.rs b/visibility-filtering/rules/user_rules.rs deleted file mode 100644 index d9ca8a87..00000000 --- a/visibility-filtering/rules/user_rules.rs +++ /dev/null @@ -1,283 +0,0 @@ -use crate::models::VfAction; -use crate::rules::{Rule, RuleContext}; -use xai_visibility_filtering::models::FilteredReason; - -#[derive(Clone)] -pub struct AuthorFlagDropRule { - name: &'static str, - flag: fn(&RuleContext<'_>) -> bool, - reason: FilteredReason, -} - -impl AuthorFlagDropRule { - pub const fn new( - name: &'static str, - flag: fn(&RuleContext<'_>) -> 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 { - if (self.flag)(context) && !context.viewer().is_author() { - return VfAction::Drop(self.reason.clone()); - } - VfAction::Allow - } -} - -pub const SUSPENDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "SuspendedAuthorRule", - |context| context.author().is_suspended(), - FilteredReason::AuthorIsSuspended, -); -pub const DEACTIVATED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DeactivatedAuthorRule", - |context| context.author().is_deactivated(), - FilteredReason::AuthorIsDeactivated, -); -pub const ERASED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "ErasedAuthorRule", - |context| context.author().is_erased(), - FilteredReason::AuthorAccountIsInactive, -); -pub const OFFBOARDED_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "OffboardedAuthorRule", - |context| context.author().is_offboarded(), - FilteredReason::AuthorAccountIsInactive, -); -pub const NSFW_USER_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DropNsfwUserAuthorRule", - |context| context.author().is_nsfw_user(), - FilteredReason::ContainNsfwMedia, -); -pub const NSFW_ADMIN_AUTHOR_DROP: AuthorFlagDropRule = AuthorFlagDropRule::new( - "DropNsfwAdminAuthorRule", - |context| context.author().is_nsfw_admin(), - FilteredReason::ContainNsfwMedia, -); - -pub struct ProtectedAuthorDropRule; - -impl Rule for ProtectedAuthorDropRule { - fn name(&self) -> &'static str { - "ProtectedAuthorDropRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - if context.author().is_protected() - && !context.viewer().is_author() - && (context.viewer().is_logged_out() || !context.viewer().follows_author()) - { - return VfAction::Drop(FilteredReason::AuthorIsProtected); - } - VfAction::Allow - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{AuthorFeatures, HydratedTweetCandidate}; - use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; - - fn candidate_with_author( - suspended: bool, - deactivated: bool, - protected: bool, - ) -> HydratedTweetCandidate { - candidate() - .with_author_features(AuthorFeatures { - is_suspended: suspended, - is_deactivated: deactivated, - is_protected: protected, - ..Default::default() - }) - .build() - } - - #[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(VIEWER_ID), &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(&author_viewer(), &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(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Drop(_) - )); - } - - #[test] - fn protected_author_allows_follower() { - let rule = ProtectedAuthorDropRule; - let mut c = candidate_with_author(false, false, true); - c.relationship.viewer_follows_author = true; - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - #[test] - fn protected_author_drops_logged_out_viewer() { - let rule = ProtectedAuthorDropRule; - let c = candidate_with_author(false, false, true); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&logged_out_viewer(), &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(&author_viewer(), &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 = candidate().with_author_features(author_features).build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &flagged)), - VfAction::Drop(FilteredReason::AuthorAccountIsInactive) - )); - let unflagged = candidate().build(); - assert!(matches!( - rule.evaluate(&crate::rules::test_context(&viewer(VIEWER_ID), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } - - fn candidate_with_nsfw_author( - is_nsfw_user: bool, - is_nsfw_admin: bool, - ) -> HydratedTweetCandidate { - candidate() - .with_author_features(AuthorFeatures { - is_nsfw_user, - is_nsfw_admin, - ..Default::default() - }) - .build() - } - - #[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(VIEWER_ID), &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(&author_viewer(), &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(VIEWER_ID), &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(&author_viewer(), &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(VIEWER_ID), &c)), - VfAction::Allow - )); - } -} From 2a38187dbffa67ae1c81781db7e048a5f461ab13 Mon Sep 17 00:00:00 2001 From: joshs1017dev Date: Tue, 1 Sep 2026 23:11:56 -0400 Subject: [PATCH 14/18] Deduplicate in_network_ids before VF lookup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit in_network_ids is passed to the VF client without deduplication, while oon_ids is deduped four lines below. retweeted_tweet_id is pushed for every candidate that has one, so the same ID repeats once per retweet of a given post — most often when that post is going viral. Neither VfClient implementation dedupes its input: StratoVfClient builds one call per element, and XaiVfClient chunks by XAI_VF_MAX_BATCH_SIZE, so duplicates consume batch slots and can force an extra round trip. Not a correctness issue — results collapse into a HashMap keyed by tweet ID — but redundant work on the For You serving path. Co-Authored-By: Claude Opus 5 --- home-mixer/candidate_hydrators/vf_candidate_hydrator.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs index 264a7a8e..b6554ed6 100644 --- a/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs +++ b/home-mixer/candidate_hydrators/vf_candidate_hydrator.rs @@ -83,6 +83,8 @@ impl Hydrator for VFCandidateHydrator { } } + in_network_ids.sort_unstable(); + in_network_ids.dedup(); oon_ids.sort_unstable(); oon_ids.dedup(); From 85ac72a1bba41f21615e3f0bca56da75970a6633 Mon Sep 17 00:00:00 2001 From: CI agent Date: Wed, 2 Sep 2026 19:46:48 +0000 Subject: [PATCH 15/18] Open-source X Recommendation Algorithm --- grox/core/data_loaders/data_types.py | 4 +- home-mixer/models/candidate.rs | 2 + home-mixer/models/query.rs | 2 + home-mixer/params/param.rs | 4 +- home-mixer/scorers/phoenix_scorer.rs | 2 + home-mixer/server.rs | 9 +- home-mixer/util/phoenix_request.rs | 1 + .../serving/xai-recsys-engine/src/python.rs | 7 +- .../xai-recsys-proto/proto/recsys.proto | 5 + .../common/xai-proto/proto/recsys.proto | 5 + .../xai_checkpointing/dek.py | 93 +++- .../xai_checkpointing/encrypted_kvstore.py | 10 +- .../xai_checkpointing/load.py | 73 ++- .../xai_checkpointing/orbax_encrypted.py | 91 ++- phoenix/xrex/configs/xrecsys_two_tower.py | 4 +- phoenix/xrex/train/checkpoint_write.py | 31 +- phoenix/xrex/train/trainer.py | 2 +- phoenix/xrex/utils/checkpointing.py | 110 ++-- scarecrow/legacy/DownstreamServices.scala | 4 + scarecrow/legacy/XReviewIntakeRegistry.scala | 179 ++++++ scarecrow/legacy/XReviewReportBuilder.scala | 59 ++ scarecrow/legacy/XReviewReportSubmitter.scala | 91 +++ visibility-filtering/filter.rs | 16 +- visibility-filtering/get_safety_labels.rs | 12 +- visibility-filtering/hydration/mod.rs | 2 +- .../hydration/safety_label_hydrator.rs | 4 +- visibility-filtering/rules/author_rules.rs | 2 +- visibility-filtering/rules/golden_corpus.rs | 12 +- visibility-filtering/rules/mod.rs | 160 +----- visibility-filtering/rules/registry.rs | 527 ++++++++++-------- visibility-filtering/rules/rule_spec.rs | 10 +- visibility-filtering/rules/tweet_rules.rs | 35 +- .../safety_label_source/source.rs | 32 +- visibility-filtering/server_deps.rs | 6 +- 34 files changed, 1081 insertions(+), 525 deletions(-) create mode 100644 scarecrow/legacy/XReviewIntakeRegistry.scala create mode 100644 scarecrow/legacy/XReviewReportBuilder.scala create mode 100644 scarecrow/legacy/XReviewReportSubmitter.scala diff --git a/grox/core/data_loaders/data_types.py b/grox/core/data_loaders/data_types.py index 0acd40eb..893a2247 100644 --- a/grox/core/data_loaders/data_types.py +++ b/grox/core/data_loaders/data_types.py @@ -373,7 +373,7 @@ def from_thrift_model( def to_convo(self) -> list[str | ConvoImage]: body: list[str | ConvoImage] = [] if self.title: - body.append(f"\n\nTitle: {self.title}") + body.append(f"\n\n[Card Title] {self.title}") if self.description: body.append(f"\n\nDescription: {self.description}") if self.domain: @@ -445,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: diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index cd728716..f94d4a4b 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -26,6 +26,8 @@ pub struct PostCandidate { 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" 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/param.rs b/home-mixer/params/param.rs index dc563640..7b3825c5 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-09-01T16:42:25Z +// mirrored from config feature-switch defaults; last sync 2026-09-02T16:03:45Z use xai_feature_switches::param; param!( @@ -134,7 +134,7 @@ param!( PhoenixRetrievalAggregationType, String, "rust_home_mixer_phoenix_retrieval_aggregation_type", - "DENSE_WITH_SHORT_DWELL" + "DENSE_WITH_LONG_DWELL" ); param!( diff --git a/home-mixer/scorers/phoenix_scorer.rs b/home-mixer/scorers/phoenix_scorer.rs index bba14c3d..e7c5aa51 100644 --- a/home-mixer/scorers/phoenix_scorer.rs +++ b/home-mixer/scorers/phoenix_scorer.rs @@ -107,6 +107,7 @@ 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), @@ -121,6 +122,7 @@ 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; diff --git a/home-mixer/server.rs b/home-mixer/server.rs index 3981dcf0..ed51d90e 100644 --- a/home-mixer/server.rs +++ b/home-mixer/server.rs @@ -329,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); @@ -461,9 +462,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 { @@ -506,6 +508,7 @@ 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; if !cursor_str.is_empty() { @@ -620,6 +623,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) @@ -720,6 +724,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/util/phoenix_request.rs b/home-mixer/util/phoenix_request.rs index ab58a3e4..16779871 100644 --- a/home-mixer/util/phoenix_request.rs +++ b/home-mixer/util/phoenix_request.rs @@ -132,6 +132,7 @@ 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(), diff --git a/phoenix/crates/serving/xai-recsys-engine/src/python.rs b/phoenix/crates/serving/xai-recsys-engine/src/python.rs index 8cd7bb47..5e0610da 100644 --- a/phoenix/crates/serving/xai-recsys-engine/src/python.rs +++ b/phoenix/crates/serving/xai-recsys-engine/src/python.rs @@ -2024,6 +2024,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()) @@ -2031,7 +2032,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]); + } } }); }); @@ -2040,6 +2044,7 @@ 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) diff --git a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index e08447ef..6bd6f661 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -114,6 +114,8 @@ message PredictNextActionsRequest { repeated ConvAssetIds conv_asset_ids = 18; PageDecodeParams pageDecode = 19; + + bool returnBackboneScores = 20; } message PageDecodeParams { @@ -256,6 +258,9 @@ message NextActionDistribution { RewardOutputs rewardOutputs = 8; SlateContext slateContext = 9; + + repeated float backboneTopLogProbs = 10; + repeated float backboneContinuousValues = 11; } message RewardOutputs { diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index e08447ef..6bd6f661 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -114,6 +114,8 @@ message PredictNextActionsRequest { repeated ConvAssetIds conv_asset_ids = 18; PageDecodeParams pageDecode = 19; + + bool returnBackboneScores = 20; } message PageDecodeParams { @@ -256,6 +258,9 @@ message NextActionDistribution { RewardOutputs rewardOutputs = 8; SlateContext slateContext = 9; + + repeated float backboneTopLogProbs = 10; + repeated float backboneContinuousValues = 11; } message RewardOutputs { diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py index 5bf88e19..4f49bab4 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py @@ -1,9 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. import json +import logging +import os import pathlib +import tempfile import time +rank_logger = logging.getLogger("rank") + TREE_DEK_NAME = "_DEK" TREE_DEK_CLAIM_NAME = "_DEK.claim" PUBLISH_TIMEOUT_SECS = 120.0 @@ -17,13 +22,15 @@ def _read_wrapped_dek(dek_path: pathlib.Path) -> dict | None: return None -def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict]: +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() @@ -34,6 +41,13 @@ def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict]: 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) @@ -44,6 +58,81 @@ def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict]: 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 = bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, dek_path)) - return raw, entry["wrapped"], entry.get("context") or {} + 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: + import xai_kms + + dek_path = path / TREE_DEK_NAME + if dek_path.exists(): + entry = _read_wrapped_dek(dek_path) or {} + raw = bytes(xai_kms.nfs.unwrap_shared_dek(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: + import xai_kms + + 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 bytes(xai_kms.nfs.unwrap_shared_dek(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"envelope at {envelope} carries a legacy per-file wrapped key, " + "not the xai-dek1 derived form" + ) + _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 index ef63e751..526edbc8 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/encrypted_kvstore.py @@ -1,10 +1,13 @@ # 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") @@ -43,7 +46,12 @@ def _envelope_spec( def encrypted_kvstore(path: pathlib.Path, kms_client, encryption_chunk_size: int) -> dict: path.mkdir(parents=True, exist_ok=True) - raw, wrapped, context = publish_tree_dek(path, kms_client) + 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 33e153f7..99de47bb 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/load.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. +import base64 import contextlib import ctypes import fcntl @@ -23,6 +24,8 @@ 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 @@ -52,35 +55,49 @@ def _names_from_tree_metadata(metadata_json: dict[str, Any]) -> list[str]: ] +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[bool, object | None, dict[str, Any], list[str], ts.Context]: - encrypted = _is_encrypted_tree(path) - if encrypted: - import xai_kms - - os.environ.setdefault("TENSORSTORE_HTTP_THREADS", "64") - if kms_client is None: - kms_client = xai_kms.KmsClient.from_cluster_env() - metadata_json = json.loads( - xai_kms.nfs.open_envelope(kms_client, str(path / "_METADATA")).read() - ) - checkpoint_names = _names_from_tree_metadata(metadata_json) - else: +) -> 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 - context_spec = { - "file_io_concurrency": {"limit": 128}, - "cache_pool#ocdbt": {"total_bytes_limit": 100000000}, - } - if encrypted: - context_spec["http_request_concurrency"] = {"limit": 128} + import xai_kms - return encrypted, kms_client, metadata_json, checkpoint_names, ts.Context(context_spec) + 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: @@ -320,7 +337,7 @@ def load_checkpoint( path = pathlib.Path(path) / tag - encrypted, kms_client, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + kvstore_base, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( path, kms_client ) use_zarr3 = metadata_json["use_zarr3"] @@ -377,10 +394,8 @@ def load_checkpoint( if node_lock is not None: stack.callback(node_lock.close) tspec_transform = None - if encrypted: - import xai_kms - - tspec_transform = stack.enter_context(xai_kms.KvServe(kms_client, path)).rewrite_ocdbt + 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(): @@ -475,7 +490,7 @@ def load_checkpoint_streamed( start = time.time() path = pathlib.Path(path) / tag - encrypted, kms_client, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( + kvstore_base, metadata_json, checkpoint_names, ts_context = _prepare_checkpoint_read( path, kms_client ) use_zarr3 = metadata_json["use_zarr3"] @@ -537,10 +552,8 @@ def load_checkpoint_streamed( if node_lock is not None: stack.callback(node_lock.close) tspec_transform = None - if encrypted: - import xai_kms - - tspec_transform = stack.enter_context(xai_kms.KvServe(kms_client, path)).rewrite_ocdbt + 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(): diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py index aa0aee45..4bcc91ef 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py @@ -29,6 +29,14 @@ def encrypt_write(base, directory, name: str, data: bytes) -> None: 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__}" @@ -80,6 +88,11 @@ def _save_fn(): 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( @@ -135,7 +148,7 @@ def close(self) -> None: return None -def _encrypted_array_handler(array_handler_cls, base_for, handler_kwargs): +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) @@ -148,6 +161,13 @@ def _get_json_tspec_write(self, info, *args, **kwargs): ) 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") @@ -165,21 +185,24 @@ async def _serialize_sharding(self, sharding, info, sharding_metadata_txn): if serialized_sharding is not None: await t.with_transaction(sharding_metadata_txn).write(serialized_sharding) - gb = handler_kwargs.get("save_concurrent_gb") - if gb is not None: - return EncryptedArrayHandler(int(gb) * 10**9) - return EncryptedArrayHandler() + return EncryptedArrayHandler(*array_handler_args) def encrypted_checkpointer( - kms_client, encryption_chunk_size: int, timeout_secs: int, array_handler_cls, handler_kwargs + 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, handler_kwargs) + 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)) ) @@ -187,17 +210,19 @@ def base_for(directory) -> dict: kms_client, encryption_chunk_size, use_ocdbt=True, - use_zarr3=handler_kwargs.get("use_zarr3", False), - save_concurrent_bytes=_pytree._concurrent_bytes(handler_kwargs.get("save_concurrent_gb")), + 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( - handler_kwargs.get("restore_concurrent_gb") + 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, **handler_kwargs + handler_impl=impl, type_handler_registry=registry, **checkpoint_handler_kwargs ), timeout_secs, checkpoint_metadata_store=EncryptedCheckpointMetadataStore(impl._base_for), @@ -226,6 +251,8 @@ def add(self, *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: @@ -240,16 +267,31 @@ def get_encrypted_checkpointer( save_concurrent_gb: int | None, array_handler_cls, ): - global _ENCRYPTED_CHECKPOINTER, _ENCRYPTED_SAVE_CONCURRENT_GB, _ENCRYPTED_TIMEOUT_SECS + global \ + _ENCRYPTED_CHECKPOINTER, \ + _ENCRYPTED_SAVE_CONCURRENT_GB, \ + _ENCRYPTED_TIMEOUT_SECS, \ + _ENCRYPTED_CHUNK_SIZE, \ + _ENCRYPTED_ARRAY_HANDLER_CLS if _ENCRYPTED_CHECKPOINTER is None: - handler_kwargs: dict = {"use_zarr3": True} + checkpoint_handler_kwargs: dict = {"use_zarr3": True} if save_concurrent_gb is not None: - handler_kwargs["save_concurrent_gb"] = save_concurrent_gb - handler_kwargs["restore_concurrent_gb"] = save_concurrent_gb + 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, handler_kwargs + 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) @@ -267,6 +309,23 @@ def get_encrypted_checkpointer( 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 diff --git a/phoenix/xrex/configs/xrecsys_two_tower.py b/phoenix/xrex/configs/xrecsys_two_tower.py index b35329d7..6fd9cb37 100644 --- a/phoenix/xrex/configs/xrecsys_two_tower.py +++ b/phoenix/xrex/configs/xrecsys_two_tower.py @@ -270,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, diff --git a/phoenix/xrex/train/checkpoint_write.py b/phoenix/xrex/train/checkpoint_write.py index f5b0f3e8..4e4232b4 100644 --- a/phoenix/xrex/train/checkpoint_write.py +++ b/phoenix/xrex/train/checkpoint_write.py @@ -66,10 +66,33 @@ def save_checkpoint( path = self.get_checkpoint_path(self.ctx) - if getattr(self, "_restored_encrypted", False): + encrypt_checkpoint = self.checkpoint_config.encrypt + if getattr(self, "_restored_encrypted", False) and not encrypt_checkpoint: raise ValueError( - "restored from an ENCRYPTED checkpoint but the orbax writer cannot " - "produce encrypted saves — refusing to downgrade to plaintext saves" + "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"): @@ -133,6 +156,8 @@ def _callback( 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/trainer.py b/phoenix/xrex/train/trainer.py index 5c74d583..12245a50 100644 --- a/phoenix/xrex/train/trainer.py +++ b/phoenix/xrex/train/trainer.py @@ -1167,7 +1167,7 @@ def warm_start_staging_spec(self): return (lambda tree: tree.purge_opt_state()), keep_fields def _uses_tensorstore_save(self) -> bool: - return self.checkpoint_config.encrypt or self.checkpoint_config.save_method == "tensorstore" + return self.checkpoint_config.save_method == "tensorstore" def maybe_load_checkpoint( self, ctx: TrainerContext, tag: str | None = None diff --git a/phoenix/xrex/utils/checkpointing.py b/phoenix/xrex/utils/checkpointing.py index 12ad67a3..953ab1a5 100644 --- a/phoenix/xrex/utils/checkpointing.py +++ b/phoenix/xrex/utils/checkpointing.py @@ -6,6 +6,7 @@ import gc import logging import os +import sys import time from typing import Any @@ -74,6 +75,14 @@ def _release_batch_memory(): 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: @@ -90,7 +99,7 @@ def get_checkpointer(timeout_secs=900, save_concurrent_gb: int | None = None): save_concurrent_gb, save_concurrent_gb, ) - _CHECKPOINTER = ocp.AsyncCheckpointer( + _CHECKPOINTER = _AsyncCheckpointer( ocp.PyTreeCheckpointHandler(**handler_kwargs), timeout_secs ) if not hasattr(_CHECKPOINTER, "_post_finalization_callback"): @@ -222,6 +231,8 @@ def save_checkpoint( 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() @@ -230,7 +241,24 @@ def save_checkpoint( tag = "orbax-ckpt" os.makedirs(path, exist_ok=True) - checkpointer = get_checkpointer(timeout_secs, save_concurrent_gb=save_concurrent_gb) + 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"): checkpointer.wait_until_finished() @@ -256,34 +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 save_concurrent_gb is not None: - concurrent_bytes = int(save_concurrent_gb) * 10**9 - if compressed: - handler = ThrottledD2HArrayHandler(concurrent_bytes) + 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: - 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, - ) + 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) @@ -299,13 +328,16 @@ 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 (save_concurrent_gb=%s, blocking=%s)", @@ -329,4 +361,8 @@ def _callback(): 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/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/visibility-filtering/filter.rs b/visibility-filtering/filter.rs index 5be5b357..032b56a1 100644 --- a/visibility-filtering/filter.rs +++ b/visibility-filtering/filter.rs @@ -1,7 +1,7 @@ use crate::hydration::{HydrationOutput, HydrationPipeline, HydrationRequest}; use crate::models::{RawCandidate, TweetId, VfAction}; 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; @@ -32,14 +32,14 @@ pub struct FilterResponse { 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, } } @@ -71,7 +71,7 @@ impl FilterTweets { .map(|candidate| { ( TweetId(candidate.tweet_id), - self.policies + self.rule_engine .evaluate(request.safety_level, &viewer_features, candidate), ) }) @@ -96,7 +96,9 @@ impl FilterTweets { 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(); @@ -192,7 +194,7 @@ pub(crate) mod test_support { labels, crate::hydration::FallbackCacheMode::Disabled, ), - Policies::new(), + RuleEngine::new(), ) } } diff --git a/visibility-filtering/get_safety_labels.rs b/visibility-filtering/get_safety_labels.rs index 84054fe4..f9630b40 100644 --- a/visibility-filtering/get_safety_labels.rs +++ b/visibility-filtering/get_safety_labels.rs @@ -87,7 +87,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 +97,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); @@ -182,8 +182,8 @@ 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, ) @@ -203,7 +203,7 @@ mod tests { 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")), @@ -261,7 +261,7 @@ mod tests { #[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/mod.rs b/visibility-filtering/hydration/mod.rs index 33e088c0..1e716aa3 100644 --- a/visibility-filtering/hydration/mod.rs +++ b/visibility-filtering/hydration/mod.rs @@ -126,7 +126,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 { 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/rules/author_rules.rs b/visibility-filtering/rules/author_rules.rs index 758ca0bd..51ff97b4 100644 --- a/visibility-filtering/rules/author_rules.rs +++ b/visibility-filtering/rules/author_rules.rs @@ -173,7 +173,7 @@ mod tests { AuthorFeatures, HydratedTweetCandidate, VfAction, ViewerAuthorRelationship, ViewerFeatures, }; use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; - use crate::rules::{test_context, Rule}; + use crate::rules::test_context; fn assert_drops( spec: &RuleSpec, diff --git a/visibility-filtering/rules/golden_corpus.rs b/visibility-filtering/rules/golden_corpus.rs index dbc3f31c..d97ed59e 100644 --- a/visibility-filtering/rules/golden_corpus.rs +++ b/visibility-filtering/rules/golden_corpus.rs @@ -5,7 +5,7 @@ use crate::models::{ use crate::rules::fixtures::{ author_viewer, candidate, logged_out_viewer, sensitive_opt_in_viewer, viewer, VIEWER_ID, }; -use crate::rules::{Policies, SafetyLevel}; +use crate::rules::{RuleEngine, SafetyLevel}; use std::collections::BTreeSet; use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; use xai_visibility_filtering::models::{ @@ -26,10 +26,10 @@ struct Case { #[test] fn golden_corpus_pins_policy_verdicts() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let mut failures = Vec::new(); for case in cases() { - let verdict = policies.evaluate(case.level, &case.viewer, &case.candidate); + 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 { @@ -54,10 +54,10 @@ fn golden_corpus_pins_policy_verdicts() { #[test] fn every_wired_rule_decides_a_corpus_case() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let wired: BTreeSet<&'static str> = [FilterAll, TimelineHome, TimelineHomeRecommendations] .into_iter() - .flat_map(|level| policies.wired_rule_names(level)) + .flat_map(|level| rule_engine.wired_rule_names(level)) .collect(); let deciders: BTreeSet<&'static str> = cases() .iter() @@ -66,7 +66,7 @@ fn every_wired_rule_decides_a_corpus_case() { let missing: Vec<&&'static str> = wired.difference(&deciders).collect(); assert!( missing.is_empty(), - "rules wired in Policies but never the decider of any corpus case: {missing:?}" + "rules wired in RuleEngine but never the decider of any corpus case: {missing:?}" ); } diff --git a/visibility-filtering/rules/mod.rs b/visibility-filtering/rules/mod.rs index fa5ab79a..d28d532b 100644 --- a/visibility-filtering/rules/mod.rs +++ b/visibility-filtering/rules/mod.rs @@ -13,12 +13,7 @@ use crate::models::VfAction; use xai_visibility_filtering::models::FilteredReason; pub use context::RuleContext; -pub use registry::{Policies, SafetyLevel}; - -pub trait Rule: Send + Sync { - fn name(&self) -> &'static str; - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction; -} +pub use registry::{RuleEngine, SafetyLevel}; #[derive(Clone, Debug)] pub struct Verdict { @@ -35,32 +30,6 @@ 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 crate::models::ViewerFeatures, @@ -77,130 +46,3 @@ pub(crate) fn test_context<'a>( &NSFW_GATING_COUNTRIES, ) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::models::{HydratedTweetCandidate, ViewerFeatures}; - 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()) - } - - #[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); - } -} diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index fdeb40b9..41a332bd 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -1,8 +1,8 @@ -use crate::models::{HydratedTweetCandidate, ViewerFeatures}; +use crate::models::{HydratedTweetCandidate, VfAction, ViewerFeatures}; use crate::params::NsfwGatingCountries; use crate::rules::rule_spec::RuleSpec; use crate::rules::{author_rules, tweet_rules}; -use crate::rules::{evaluate_rules, Rule, RuleContext, Verdict}; +use crate::rules::{RuleContext, Verdict}; use std::sync::Arc; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -22,32 +22,108 @@ 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 Policy { + pub(super) const fn new(rules: &'static [&'static [RuleSpec]]) -> Self { + Self { + rules, + additional_rules: &[], + } + } + + fn rules(&self) -> impl Iterator + '_ { + self.rules + .iter() + .chain(self.additional_rules) + .copied() + .flatten() + } + + 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(_) => {} + } + } + + Verdict { action, decided_by } + } + + fn len(&self) -> usize { + self.rules().count() + } + + #[cfg(test)] + fn rule_names(&self) -> impl Iterator + '_ { + self.rules().map(RuleSpec::name) + } +} + +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, +}; + +pub struct RuleEngine { nsfw_gating_countries: Arc, } -impl Policies { +impl RuleEngine { pub fn new() -> Self { Self::with_nsfw_gating_countries(Arc::new(NsfwGatingCountries::new())) } pub fn with_nsfw_gating_countries(gating_countries: Arc) -> Self { Self { - filter_all: rule_specs(tweet_rules::FILTER_ALL).collect(), - timeline_home: timeline_home_policy(), - timeline_home_recommendations: timeline_home_recommendations_policy(), nsfw_gating_countries: gating_countries, } } - fn select(&self, level: SafetyLevel) -> &[Box] { + fn select(level: SafetyLevel) -> &'static Policy { match level { - SafetyLevel::FilterAll => &self.filter_all, - SafetyLevel::TimelineHome => &self.timeline_home, - SafetyLevel::TimelineHomeRecommendations => &self.timeline_home_recommendations, + SafetyLevel::FilterAll => &FILTER_ALL_POLICY, + SafetyLevel::TimelineHome => &TIMELINE_HOME_POLICY, + SafetyLevel::TimelineHomeRecommendations => &TIMELINE_HOME_RECOMMENDATIONS_POLICY, } } @@ -58,62 +134,28 @@ impl Policies { candidate: &HydratedTweetCandidate, ) -> Verdict { let context = RuleContext::new(level, viewer, candidate, &self.nsfw_gating_countries); - evaluate_rules(self.select(level), &context) + Self::select(level).evaluate(&context) } #[cfg(test)] pub(crate) fn wired_rule_names(&self, level: SafetyLevel) -> Vec<&'static str> { - self.select(level).iter().map(|rule| rule.name()).collect() + Self::select(level).rule_names().collect() } pub fn rule_counts(&self) -> (usize, usize) { ( - self.timeline_home.len(), - self.timeline_home_recommendations.len(), + TIMELINE_HOME_POLICY.len(), + TIMELINE_HOME_RECOMMENDATIONS_POLICY.len(), ) } } -impl Default for Policies { +impl Default for RuleEngine { fn default() -> Self { Self::new() } } -fn rule_specs(specs: &'static [RuleSpec]) -> impl Iterator> { - specs - .iter() - .map(|spec| Box::new(spec.clone()) as Box) -} - -fn base_home_rules() -> Vec> { - let mut rules: Vec> = Vec::new(); - rules.extend(rule_specs(author_rules::AUTHOR_STATE_DROPS)); - rules.extend(rule_specs(author_rules::SOCIALGRAPH_DROPS)); - rules.extend(rule_specs(tweet_rules::TWEET_LABEL_DROPS)); - rules.extend(rule_specs(tweet_rules::NULLCAST_DROP)); - rules.extend(rule_specs(tweet_rules::TES_HOME_DROPS)); - rules.extend(rule_specs(tweet_rules::SENSITIVE_VIEWER_DROPS)); - rules.extend(rule_specs(tweet_rules::EXCLUSIVE_TWEET_DROP)); - rules.extend(rule_specs(tweet_rules::NSFW_MEDIA_INTERSTITIALS)); - rules.extend(rule_specs(tweet_rules::NSFW_AUTHOR_INTERSTITIAL)); - rules -} - -fn timeline_home_policy() -> Vec> { - base_home_rules() -} - -fn timeline_home_recommendations_policy() -> Vec> { - let mut rules = base_home_rules(); - rules.extend(rule_specs(tweet_rules::RECS_MEDIA_DROPS)); - rules.extend(rule_specs(author_rules::OON_NSFW_AUTHOR_DROPS)); - rules.extend(rule_specs(tweet_rules::OON_TWEET_FLAG_DROPS)); - rules.extend(rule_specs(tweet_rules::OON_TWEET_LABEL_DROPS)); - rules.extend(rule_specs(author_rules::OON_USER_LABEL_DROPS)); - rules -} - #[cfg(test)] mod tests { use super::*; @@ -121,58 +163,11 @@ mod tests { HydratedTweetCandidate, MediaFeature, TweetFeatures, VfAction, ViewerFeatures, }; use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; - use xai_visibility_filtering::models::FilteredReason; - - struct RecommendationsOnlyRule; - - impl Rule for RecommendationsOnlyRule { - fn name(&self) -> &'static str { - "RecommendationsOnlyRule" - } - - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { - match context.safety_level() { - SafetyLevel::TimelineHomeRecommendations => { - VfAction::Drop(FilteredReason::UnspecifiedReason) - } - SafetyLevel::FilterAll | SafetyLevel::TimelineHome => VfAction::Allow, - } - } - } - - #[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)], - nsfw_gating_countries: Arc::new(NsfwGatingCountries::new()), - }; - 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(_) - )); - } #[test] fn refreshed_config_country_reaches_the_wired_rule() { let gating_countries = Arc::new(NsfwGatingCountries::new()); - let policies = Policies::with_nsfw_gating_countries(Arc::clone(&gating_countries)); + 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() @@ -183,7 +178,7 @@ mod tests { ..viewer(VIEWER_ID) }; - let verdict = policies.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); + let verdict = rule_engine.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); assert!(!matches!(verdict.action, VfAction::Drop(_))); gating_countries.refresh_from( @@ -199,7 +194,7 @@ rust_vf: ) .unwrap(), ); - let verdict = policies.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); + let verdict = rule_engine.evaluate(SafetyLevel::TimelineHome, &viewer, &candidate); assert!(matches!(verdict.action, VfAction::Drop(_))); assert_eq!( verdict.decided_by, @@ -207,25 +202,14 @@ rust_vf: ); } - #[test] - fn filter_all_rule_drops_even_self_view() { - let candidate = candidate().build(); - let viewer = author_viewer(); - let spec = &tweet_rules::FILTER_ALL[0]; - assert!(matches!( - spec.evaluate(&crate::rules::test_context(&viewer, &candidate)), - VfAction::Drop(_) - )); - } - #[test] fn wired_rule_order_matches_pre_migration_sequence() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); assert_eq!( - policies.wired_rule_names(SafetyLevel::FilterAll), + rule_engine.wired_rule_names(SafetyLevel::FilterAll), vec!["FilterAllRule"] ); - let home = policies.wired_rule_names(SafetyLevel::TimelineHome); + let home = rule_engine.wired_rule_names(SafetyLevel::TimelineHome); assert_eq!( home, vec![ @@ -289,33 +273,40 @@ rust_vf: "DoNotAmplifyNonFollowerRule", ]); assert_eq!( - policies.wired_rule_names(SafetyLevel::TimelineHomeRecommendations), + rule_engine.wired_rule_names(SafetyLevel::TimelineHomeRecommendations), recs ); } #[test] fn filter_all_policy_drops_pristine_candidate() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate().build(); - let verdict = policies.evaluate( + let verdict = rule_engine.evaluate( SafetyLevel::FilterAll, &ViewerFeatures::default(), &candidate, ); assert!(matches!(verdict.action, VfAction::Drop(_))); - let verdict = policies.evaluate( + let verdict = rule_engine.evaluate( SafetyLevel::TimelineHome, &ViewerFeatures::default(), &candidate, ); assert!(matches!(verdict.action, VfAction::Allow)); + + let verdict = rule_engine.evaluate( + SafetyLevel::TimelineHomeRecommendations, + &ViewerFeatures::default(), + &candidate, + ); + assert!(matches!(verdict.action, VfAction::Allow)); } #[test] fn dmca_media_drops_recommendations_only() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate() .with_tweet_features(TweetFeatures { media: MediaFeature { @@ -326,14 +317,14 @@ rust_vf: }) .build(); - let timeline_home = policies.evaluate( + let timeline_home = rule_engine.evaluate( SafetyLevel::TimelineHome, &ViewerFeatures::default(), &candidate, ); assert!(matches!(timeline_home.action, VfAction::Allow)); - let recommendations = policies.evaluate( + let recommendations = rule_engine.evaluate( SafetyLevel::TimelineHomeRecommendations, &ViewerFeatures::default(), &candidate, @@ -344,7 +335,7 @@ rust_vf: #[test] fn tweet_nsfw_flag_drops_recommendations_only() { use crate::models::NsfwFeature; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate() .with_tweet_features(TweetFeatures { nsfw: NsfwFeature { @@ -356,7 +347,7 @@ rust_vf: .build(); let viewer = viewer(VIEWER_ID); - let timeline_home = policies + let timeline_home = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -364,7 +355,7 @@ rust_vf: "in-network tweet nsfw_user flag should allow, got {timeline_home:?}" ); - let recommendations = policies.evaluate( + let recommendations = rule_engine.evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, &candidate, @@ -376,7 +367,7 @@ rust_vf: #[test] fn nsfw_author_interstitials_in_network_but_drops_oon() { use crate::models::AuthorFeatures; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate() .with_media() .with_author_features(AuthorFeatures { @@ -386,7 +377,7 @@ rust_vf: .build(); let viewer = viewer(VIEWER_ID); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -394,7 +385,7 @@ rust_vf: "in-network NSFW author should interstitial, got {in_network:?}" ); - let oon = policies + let oon = rule_engine .evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, @@ -411,7 +402,7 @@ rust_vf: fn egregious_nsfw_does_not_drop() { use crate::models::SafetyLabelType; use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let tweet_candidate = candidate() .with_label(SafetyLabelType::EGREGIOUS_NSFW) @@ -420,14 +411,14 @@ rust_vf: let viewer = viewer(VIEWER_ID); for candidate in [&tweet_candidate, &user_candidate] { - let in_network = policies + let in_network = rule_engine .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 + let oon = rule_engine .evaluate(SafetyLevel::TimelineHomeRecommendations, &viewer, candidate) .action; assert!( @@ -446,67 +437,16 @@ rust_vf: c } - #[test] - fn fosnr_labels_drop_non_author_non_follower_on_both_surfaces() { - use crate::models::SafetyLabelType; - let policies = Policies::new(); - let viewer = viewer(VIEWER_ID); - 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:?}" - ); - } - } - } - - #[test] - fn fosnr_never_drops_author() { - use crate::models::SafetyLabelType; - let policies = Policies::new(); - let author = author_viewer(); - 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:?}" - ); - } - } - } - #[test] fn fosnr_abuse_insults_drops_oon_but_allows_in_network() { use crate::models::SafetyLabelType; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let viewer = viewer(VIEWER_ID); let author = author_viewer(); for follows in [true, false] { let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, follows); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -516,7 +456,7 @@ rust_vf: } let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, false); - let oon = policies + let oon = rule_engine .evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, @@ -528,7 +468,7 @@ rust_vf: "OON FosnrAbuseInsults should drop non-author, got {oon:?}" ); - let oon_author = policies + let oon_author = rule_engine .evaluate( SafetyLevel::TimelineHomeRecommendations, &author, @@ -543,7 +483,7 @@ rust_vf: #[test] fn geo_restricted_media_drops_oon_but_allows_in_network() { - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate() .with_tweet_features(TweetFeatures { media: MediaFeature { @@ -558,7 +498,7 @@ rust_vf: ..viewer(VIEWER_ID) }; - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -566,7 +506,7 @@ rust_vf: "in-network geo-restricted media should allow (Scala wires the rule in THR only), got {in_network:?}" ); - let oon = policies.evaluate( + let oon = rule_engine.evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, &candidate, @@ -582,11 +522,11 @@ rust_vf: #[test] fn nsfw_text_drops_oon_but_allows_in_network() { use crate::models::SafetyLabelType; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate().with_label(SafetyLabelType::NSFW_TEXT).build(); let viewer = viewer(VIEWER_ID); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -594,7 +534,7 @@ rust_vf: "in-network NsfwText should allow (Scala drops it OON only), got {in_network:?}" ); - let oon = policies + let oon = rule_engine .evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, @@ -619,11 +559,11 @@ rust_vf: #[test] fn nsfw_avatar_user_label_drops_oon_but_allows_in_network() { use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate_with_author_user_label(LabelValue::NSFW_AVATAR_IMAGE, false); let viewer = viewer(VIEWER_ID); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -631,7 +571,7 @@ rust_vf: "in-network NsfwAvatarImage should allow, got {in_network:?}" ); - let oon = policies.evaluate( + let oon = rule_engine.evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, &candidate, @@ -647,12 +587,12 @@ rust_vf: #[test] fn recommendations_blacklist_does_not_drop() { use xai_x_thrift::user_labels::LabelValue; - let policies = Policies::new(); + let rule_engine = RuleEngine::new(); let candidate = candidate_with_author_user_label(LabelValue::RECOMMENDATIONS_BLACKLIST, false); let viewer = viewer(VIEWER_ID); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -660,7 +600,7 @@ rust_vf: "in-network RecommendationsBlacklist should allow, got {in_network:?}" ); - let oon = policies + let oon = rule_engine .evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, @@ -673,33 +613,14 @@ rust_vf: ); } - #[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 = viewer(VIEWER_ID); - - 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:?}" - ); - } - } - #[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 rule_engine = RuleEngine::new(); let candidate = candidate_with_author_user_label(LabelValue::ABUSIVE_HIGH_RECALL, false); let viewer = viewer(VIEWER_ID); - let in_network = policies + let in_network = rule_engine .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) .action; assert!( @@ -707,7 +628,7 @@ rust_vf: "in-network AbusiveHighRecall should allow, got {in_network:?}" ); - let oon = policies.evaluate( + let oon = rule_engine.evaluate( SafetyLevel::TimelineHomeRecommendations, &viewer, &candidate, @@ -719,4 +640,166 @@ rust_vf: ); assert_eq!(oon.decided_by, Some("AbusiveHighRecallRule")); } + + #[derive(Debug, PartialEq, Eq)] + enum RowClass { + Drop, + Interstitial, + } + + 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, + } + } + + 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 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() + ); + } + } + + 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 } + } + + fn allow(_: &RuleContext<'_>) -> VfAction { + VfAction::Allow + } + + fn drop_suspended(_: &RuleContext<'_>) -> VfAction { + VfAction::Drop(FilteredReason::AuthorIsSuspended) + } + + fn interstitial_nsfw(_: &RuleContext<'_>) -> VfAction { + VfAction::Interstitial(FilteredReason::ContainNsfwMedia) + } + + fn interstitial_unspecified(_: &RuleContext<'_>) -> VfAction { + VfAction::Interstitial(FilteredReason::UnspecifiedReason) + } + + fn unreachable_after_drop(_: &RuleContext<'_>) -> VfAction { + panic!("a rule after a Drop must never be evaluated"); + } + + fn context_inputs() -> (ViewerFeatures, HydratedTweetCandidate) { + (ViewerFeatures::default(), HydratedTweetCandidate::default()) + } + + 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")); + } + + #[test] + fn first_interstitial_sticks_without_short_circuit() { + let (viewer, candidate) = context_inputs(); + + let verdict = INTERSTITIALS.evaluate(&test_context(&viewer, &candidate)); + + assert!(matches!( + verdict.action, + VfAction::Interstitial(FilteredReason::ContainNsfwMedia) + )); + assert_eq!(verdict.decided_by, Some("first_interstitial")); + } + + #[test] + fn drop_after_interstitial_wins() { + let (viewer, candidate) = context_inputs(); + + let verdict = DROP_AFTER_INTERSTITIAL.evaluate(&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 (viewer, candidate) = context_inputs(); + + let verdict = ALL_ALLOWS.evaluate(&test_context(&viewer, &candidate)); + + 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 index b43c1987..6bac699b 100644 --- a/visibility-filtering/rules/rule_spec.rs +++ b/visibility-filtering/rules/rule_spec.rs @@ -1,9 +1,8 @@ use crate::models::VfAction; use crate::rules::context::{AuthorPredicates, TweetPredicates}; -use crate::rules::{Rule, RuleContext}; +use crate::rules::RuleContext; use xai_visibility_filtering::models::FilteredReason; -#[derive(Clone)] pub(super) enum RuleSpec { Tweet { name: &'static str, @@ -23,14 +22,13 @@ pub(super) enum RuleSpec { }, } -#[derive(Clone)] pub(super) enum RuleAction { Drop(FilteredReason), SensitiveMediaInterstitial(FilteredReason), } -impl Rule for RuleSpec { - fn name(&self) -> &'static str { +impl RuleSpec { + pub(super) fn name(&self) -> &'static str { match self { RuleSpec::Tweet { name, .. } | RuleSpec::Author { name, .. } @@ -38,7 +36,7 @@ impl Rule for RuleSpec { } } - fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { + pub(super) fn evaluate(&self, context: &RuleContext<'_>) -> VfAction { match self { RuleSpec::Tweet { when, diff --git a/visibility-filtering/rules/tweet_rules.rs b/visibility-filtering/rules/tweet_rules.rs index 7a104be3..d868a5eb 100644 --- a/visibility-filtering/rules/tweet_rules.rs +++ b/visibility-filtering/rules/tweet_rules.rs @@ -340,7 +340,7 @@ mod tests { use crate::rules::fixtures::{ author_viewer, candidate, logged_out_viewer, sensitive_opt_in_viewer, viewer, VIEWER_ID, }; - use crate::rules::{test_context, Rule, RuleContext}; + use crate::rules::{test_context, RuleContext}; use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; fn assert_drops( @@ -1122,4 +1122,37 @@ mod tests { ); assert_allows(&spec, &author_viewer(), &candidate().build()); } + + 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/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/server_deps.rs b/visibility-filtering/server_deps.rs index d59081d4..4e7501af 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -214,9 +214,9 @@ pub async fn build_prod_server( 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 policies = crate::rules::Policies::with_nsfw_gating_countries(gating_countries); - let (home_rule_count, recommendations_rule_count) = policies.rule_counts(); - let filter_tweets = FilterTweets::new(hydration_pipeline, policies); + 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; From 9b0dc319691b76088266d0d2b48faf22d2b8a82a Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 4 Sep 2026 00:37:40 +0000 Subject: [PATCH 16/18] Open-source X Recommendation Algorithm --- .../service-lib/rules/enforcement_user.yaml | 7 + .../service-lib/src/lib.rs | 1 + .../service-lib/src/rules.rs | 49 +++ grox/config/config.py | 1 + grox/flows/ptos/classifier.py | 28 +- grox/flows/ptos/task_safety_ptos_policy.py | 5 +- .../classifier_multi_step_reply_spam.py | 226 ++++++++++ grox/flows/reply_spam/generators.py | 8 +- .../reply_spam/plan_multi_step_reply_spam.py | 29 ++ grox/flows/reply_spam/prompts.py | 5 + .../reply_spam/state_multi_step_reply_spam.py | 13 + grox/flows/reply_spam/task_filter.py | 2 +- .../reply_spam/task_multi_step_reply_spam.py | 206 ++++++++++ .../filters/brazil_2026_election_filter.rs | 343 +++++++++++++++- home-mixer/params/param.rs | 8 +- home-mixer/server.rs | 24 +- home-mixer/util/phoenix_request.rs | 15 +- home-mixer/util/strato_context.rs | 47 ++- phoenix/crates/common/xai-recsys/src/util.rs | 146 ++++++- .../serving/xai-recsys-engine/Cargo.toml | 1 - .../serving/xai-recsys-engine/pyproject.toml | 1 - .../xai-recsys-engine/src/checkpoint_proxy.rs | 111 ++++- .../xai-recsys-engine/src/emb_table.rs | 19 +- .../serving/xai-recsys-engine/src/lib.rs | 1 - .../serving/xai-recsys-engine/src/python.rs | 9 - .../xai-recsys-engine/src/sid_client.rs | 189 --------- .../xai-recsys-engine/xai_recsys_engine.pyi | 5 - phoenix/xrex/cutedsl/ranker_attention_fa4.py | 214 ++++++---- phoenix/xrex/inference/launch_inference.py | 15 +- phoenix/xrex/inference/model_runner.py | 25 +- .../xrex/inference/sid_retrieval_runner.py | 17 - phoenix/xrex/models/recsys_attention.py | 22 +- phoenix/xrex/models/remat.py | 17 + .../clients/socialgraph_client.rs | 42 +- visibility-filtering/config.rs | 125 +----- visibility-filtering/dark_traffic_setup.rs | 1 + visibility-filtering/filter.rs | 56 +-- visibility-filtering/filter_tweets.rs | 82 +--- visibility-filtering/get_safety_labels.rs | 40 -- visibility-filtering/hydration/batch.rs | 19 - .../hydration/fallback_cache.rs | 112 ++--- .../hydration/gizmoduck_hydrator.rs | 78 ++-- visibility-filtering/hydration/metrics.rs | 3 - visibility-filtering/hydration/mod.rs | 101 +---- .../hydration/socialgraph_hydrator.rs | 14 +- .../hydration/tes_hydrator.rs | 36 +- visibility-filtering/lib.rs | 24 ++ visibility-filtering/main.rs | 12 + visibility-filtering/models/mod.rs | 4 +- visibility-filtering/models/safety_labels.rs | 32 +- visibility-filtering/models/tweet.rs | 9 +- visibility-filtering/params.rs | 13 +- visibility-filtering/reference_compare.rs | 16 +- visibility-filtering/rules/author_rules.rs | 30 +- visibility-filtering/rules/context.rs | 24 +- visibility-filtering/rules/fixtures.rs | 87 +++- visibility-filtering/rules/golden_corpus.rs | 107 ++++- visibility-filtering/rules/mod.rs | 9 +- visibility-filtering/rules/registry.rs | 386 +----------------- visibility-filtering/rules/tweet_rules.rs | 179 +++----- .../safety_label_source/codec.rs | 221 +--------- .../safety_label_source/lookup.rs | 4 +- .../safety_label_source/manhattan.rs | 10 - .../safety_label_source/metrics.rs | 2 - .../safety_label_source/mod.rs | 2 +- .../safety_label_source/twemcache.rs | 4 +- .../safety_label_source/types.rs | 2 - .../safety_label_source/warmer.rs | 33 +- visibility-filtering/server_deps.rs | 63 ++- visibility-filtering/twemcache/connection.rs | 25 +- visibility-filtering/twemcache/host_pool.rs | 37 +- visibility-filtering/twemcache/key.rs | 17 +- visibility-filtering/twemcache/ring.rs | 10 +- 73 files changed, 1954 insertions(+), 1926 deletions(-) create mode 100644 grox/flows/reply_spam/classifier_multi_step_reply_spam.py create mode 100644 grox/flows/reply_spam/plan_multi_step_reply_spam.py create mode 100644 grox/flows/reply_spam/state_multi_step_reply_spam.py create mode 100644 grox/flows/reply_spam/task_multi_step_reply_spam.py delete mode 100644 phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs diff --git a/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml b/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml index 30a8b0ba..b58d5a31 100644 --- a/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml +++ b/abuse-enforcement-service/service-lib/rules/enforcement_user.yaml @@ -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 diff --git a/abuse-enforcement-service/service-lib/src/lib.rs b/abuse-enforcement-service/service-lib/src/lib.rs index 991a4265..89cd1f11 100644 --- a/abuse-enforcement-service/service-lib/src/lib.rs +++ b/abuse-enforcement-service/service-lib/src/lib.rs @@ -1964,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", diff --git a/abuse-enforcement-service/service-lib/src/rules.rs b/abuse-enforcement-service/service-lib/src/rules.rs index 2dd11592..ef081aa5 100644 --- a/abuse-enforcement-service/service-lib/src/rules.rs +++ b/abuse-enforcement-service/service-lib/src/rules.rs @@ -670,6 +670,7 @@ mod tests { for id in [ "user_in_allowlist", "user_not_found", + "very_high_follower_count", "high_follower_count", "pagerank_skipped", ] { @@ -679,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", @@ -1270,6 +1282,43 @@ rules: ); } + #[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); diff --git a/grox/config/config.py b/grox/config/config.py index 3fd6a690..8a030c81 100644 --- a/grox/config/config.py +++ b/grox/config/config.py @@ -86,6 +86,7 @@ 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_X_ALGO = "eapi-grok-4-5-x-algo" diff --git a/grox/flows/ptos/classifier.py b/grox/flows/ptos/classifier.py index 660491db..554d5964 100644 --- a/grox/flows/ptos/classifier.py +++ b/grox/flows/ptos/classifier.py @@ -200,6 +200,8 @@ class SafetyPtosPolicyCrossValidator: 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) @@ -351,17 +353,16 @@ async def _validate_violent_media( 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_5" + 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: - 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 - ) + 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( @@ -370,7 +371,7 @@ async def _validate_illegal_and_regulated_behaviors( ) Metrics.counter(metric).add(1, attributes={"outcome": "unparseable"}) return _policy_no_violation( - "illegal_and_regulated_behaviors_grok_4_5_parse_error" + "illegal_and_regulated_behaviors_grok_4_1_parse_error" ) if confirm.policyType == policy.policyType: logger.info( @@ -389,7 +390,7 @@ async def _validate_illegal_and_regulated_behaviors( Metrics.counter(metric).add(1, attributes={"outcome": "disagreed"}) return _policy_no_violation( confirm.reason - or f"illegal_and_regulated_behaviors_grok_4_5_disagreed: cv={confirm.policyType.value}" + or f"illegal_and_regulated_behaviors_grok_4_1_disagreed: cv={confirm.policyType.value}" ) except Exception: logger.error( @@ -398,7 +399,7 @@ async def _validate_illegal_and_regulated_behaviors( ) Metrics.counter(metric).add(1, attributes={"outcome": "error"}) return _policy_no_violation( - "illegal_and_regulated_behaviors_grok_4_5_sample_error" + "illegal_and_regulated_behaviors_grok_4_1_sample_error" ) @@ -487,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( @@ -586,6 +586,7 @@ def build_convo( USE_GEMMA_CATEGORIES = { SafetyPolicyCategory.Spam, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, } USE_THREAD_RENDERER_CATEGORIES = { @@ -708,19 +709,14 @@ async def _sample_4_6_internal(self, convo: Conversation) -> str: ) 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 diff --git a/grox/flows/ptos/task_safety_ptos_policy.py b/grox/flows/ptos/task_safety_ptos_policy.py index efce237a..d0d7ea13 100644 --- a/grox/flows/ptos/task_safety_ptos_policy.py +++ b/grox/flows/ptos/task_safety_ptos_policy.py @@ -97,7 +97,10 @@ async def _exec_with_post(cls, ctx: TaskContext, post: Post) -> None: violation.safetyPolicy = await cls.cross_validator.validate( violation.category, post, policy ) - elif violation.category == SafetyPolicyCategory.ViolentMedia: + elif violation.category in ( + SafetyPolicyCategory.ViolentMedia, + SafetyPolicyCategory.IllegalAndRegulatedBehaviors, + ): policy = await active_classifier.classify_policy_for_violation( post, violation ) 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/generators.py b/grox/flows/reply_spam/generators.py index 5e5c080e..a8fa683a 100644 --- a/grox/flows/reply_spam/generators.py +++ b/grox/flows/reply_spam/generators.py @@ -3,6 +3,7 @@ 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 ( REPLY_RANKING, @@ -19,6 +20,7 @@ class ReplyRankingTaskGenerator(StreamTaskGenerator): PlanReplyRanking.KEY, PlanSpamComment.KEY, PlanCoordinatedSpam.KEY, + PlanMultiStepReplySpam.KEY, } def _get_loader(self): @@ -28,7 +30,11 @@ def _get_loader(self): @register class ReplyRankingRecoveryTaskGenerator(StreamTaskGenerator): TASK_GENERATOR_TYPE = REPLY_RANKING_RECOVERY - PLANS_TO_INJECT = {PlanReplyRanking.KEY, PlanSpamComment.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/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/task_filter.py b/grox/flows/reply_spam/task_filter.py index 4565abfd..91467f38 100644 --- a/grox/flows/reply_spam/task_filter.py +++ b/grox/flows/reply_spam/task_filter.py @@ -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 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/home-mixer/filters/brazil_2026_election_filter.rs b/home-mixer/filters/brazil_2026_election_filter.rs index 1d0886ee..04bb9a09 100644 --- a/home-mixer/filters/brazil_2026_election_filter.rs +++ b/home-mixer/filters/brazil_2026_election_filter.rs @@ -19,6 +19,8 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // 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. @@ -42,7 +44,6 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @DANIELBRSOARES no live account found. // @DANILOBALASOFICIAL no live account found. // @DANILOTORRES100 no live account found. -// @DECIOLIMAPT no live account found. // @DELEGADOEGUCHI no live account found. // @DEMAOLIVEIRA70 no live account found. // @DEPCELSOSABINO no live account found. @@ -57,10 +58,12 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @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. @@ -68,8 +71,8 @@ use xai_candidate_pipeline::filter::{Filter, FilterResult}; // @NETOFEITOSA6891 no live account found. // @PATRICIACRIZANTO2 no live account found. // @PAULOMOURAOTO no live account found. -// @PEDRONASSIF_RJ 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. @@ -84,6 +87,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| FxHashSet::from_iter([ // @madeleinelacsko 9179462, + // @nanacachen + 13258012, // @renildo 14160928, // @renatoroseno @@ -94,6 +99,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 15585094, // @soninhafrancine 15768105, + // @AecioNeves + 15828014, // @tatyanavaleria 15908023, // @ClariceChacon @@ -124,6 +131,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 20242549, // @marcelvanhattem 21069302, + // @lgmbrasilia + 21438601, // @nilvanferreira 21571098, // @RafaellMilas @@ -150,14 +159,20 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 25164669, // @pauloteixeira13 25562342, + // @crismonteirosp + 25841252, // @ManuelaDavila 25858078, // @tomioyano 26214420, // @CintyaMuniz 26284058, + // @rodrigolimamdh + 26498680, // @Biango 26560559, + // @samiabomfim + 27703690, // @profsta 27883076, // @Jfelippeneto @@ -184,6 +199,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 32160446, // @cirogomes 33374761, + // @gabrielmagno_13 + 33717372, // @ticokuzma 33759666, // @CristinaMel @@ -202,6 +219,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 34618485, // @BetoRicha 34665220, + // @marconiperillo + 34790309, // @Donato_PT 34795040, // @jooliveirapb @@ -256,6 +275,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 38659819, // @LucasCalilGo 38929561, + // @jandira_feghali + 39116717, // @tourinhopedro 39818820, // @murilogaldinopb @@ -278,6 +299,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 41590536, // @profcassiano 42106436, + // @KleberfreireAdv + 42202762, // @RaoniMendes 42284334, // @lcbusato @@ -346,12 +369,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 46958570, // @DanielVilela15 47371488, + // @RicMarques_RM + 47387995, // @chrispuppi 47418811, // @zeca_dirceu 47461491, // @augustocury 47529845, + // @policarpodf + 47851105, // @Altineu 47991805, // @fernandojordao @@ -380,6 +407,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 49818264, // @ZimbaldiRafa 49848446, + // @acmnetoba + 49966078, // @ReginaldoLopes 50088692, // @RicardoBarrosPP @@ -450,6 +479,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 53073647, // @neyleprevost 53163365, + // @DayseHansa + 53473894, // @renanroto 53544192, // @advandrebarros @@ -458,6 +489,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 53719510, // @paulosalimmaluf 53776600, + // @tatiroque + 53911544, // @kikosilveira 53998138, // @VINIANZILIERO @@ -500,6 +533,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 55558067, // @stephanesjunior 55607796, + // @eduardo_vidal + 55686825, // @vitor_bicca 56081072, // @Jrsantosrosa @@ -584,6 +619,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 59855833, // @subgonzagamg 59868993, + // @nandopinheiro22 + 60150569, // @marcio_motta 60469505, // @anapaulagold @@ -592,10 +629,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 60731692, // @Isquierdorio 60805457, + // @andrerodini + 60879681, // @luisbenambl 60983130, // @CovattiFilho 60994156, + // @FelicioRamuth + 61119320, // @franzepiaui 61190865, // @mariocaixa @@ -650,6 +691,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 64460310, // @josenunes_ARI 64482750, + // @hmecabo + 64493005, // @helencabral13 64493500, // @romulorippa @@ -674,6 +717,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 66459281, // @miguelcoelhope 66525428, + // @Brigadeiromma + 66529851, // @VerGuilherme 66704302, // @cantojocelito @@ -746,6 +791,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 69020170, // @juliophilbert 69073488, + // @MonicaSeixas + 69370960, + // @esuplicy + 69373037, // @juniorsinforma 69874322, // @Gyselle_Soares @@ -772,6 +821,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 71098452, // @lindberghfarias 71310152, + // @FabioFayad + 71345547, // @Daniel_PCdoB 71545154, // @FaissalCalil @@ -820,6 +871,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 74738674, // @VICENTINHOPT 74762633, + // @MCrivella + 74860967, // @agenorsantospa 74867163, // @caetraven @@ -878,6 +931,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 78714361, // @profdorinha 79174387, + // @colattodeputado + 79252957, // @gilbertoabramo 80095058, // @michelschlemper @@ -1024,6 +1079,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 93894534, // @tomazteixeira 93958167, + // @DrRafaelFavatto + 94115118, // @Obsevador 94167900, // @coroneldavidms @@ -1042,6 +1099,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 95526088, // @overissimo 95621286, + // @deciolimasc + 95699837, // @anapaulalimapt 95939603, // @PROFTULIO @@ -1156,6 +1215,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 117594353, // @ronaldornrn 117801614, + // @patybenzaquem + 118492432, // @ZeRicardoAM 119079224, // @julio_cesar_pi @@ -1226,6 +1287,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 129837652, // @alexandrebaldy 130620293, + // @paulaodopt + 131004593, + // @minc_rj + 131286287, // @Miriampetrone 132187525, // @depdelmasso @@ -1256,10 +1321,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 137548563, // @tadeuveneri 137919701, + // @sigageraldoluis + 138172915, // @Eduardo_Cury 138504441, // @Katiadiasjf 139059131, + // @rodrigominotto + 139809677, // @MerlongSolano 140413929, // @raphaelsebba @@ -1312,10 +1381,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 152900822, // @CezinhaNunes 153131779, + // @AntonioTelesJr + 153268918, // @RicardoCappelli 153563550, // @JarbasFilho_ 155230905, + // @FabiolaMansur_ + 155564771, // @dacassia1 155768056, // @sgtalexandre @@ -1354,6 +1427,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 163546697, // @MariliaPFerrari 163606193, + // @Ruicpimenta29 + 163831440, // @DepZeMilton 163935204, // @wandnogueira @@ -1408,12 +1483,18 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 168653875, // @TiberioLimeira 168657354, + // @edbernasil + 168814204, // @pepecollaco 169158984, + // @DepEzequielRN + 169180911, // @matheusmanholer 170176086, // @pedrofrancez 170188044, + // @depJeferson10 + 170505916, // @walterlfcaval 170638771, // @wilsonsousajr @@ -1470,6 +1551,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 185482131, // @rodfvale 186331377, + // @Andiara1400 + 187726597, // @acarlosmendes 190308328, // @pablomarcal @@ -1480,6 +1563,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 192658045, // @pretagilsa 193045733, + // @joeldaharpa + 195423732, // @MarciaTaschetti 198338329, // @JacksonAndre7 @@ -1562,6 +1647,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 235765762, // @euberlucas 237071521, + // @lyneoliver + 237497430, // @elzefacchinetti 237562432, // @wildermorais @@ -1602,7 +1689,7 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 255300173, // @wellington_luiz 255637975, - // @depfrederico55 + // @fredantunes_rs 256499222, // @DelegadoJacovos 257247983, @@ -1630,6 +1717,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 267467458, // @alinegurgel_ap 267981392, + // @depsebastiao + 268318846, // @pastorflamarion 269618056, // @ayres_jr @@ -1648,6 +1737,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 274688443, // @netocoelhoo 276794985, + // @DafneOrion + 277586067, // @danealencar 278126758, // @CostaMarinara @@ -1722,6 +1813,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 310502840, // @TercioTinoco 311616488, + // @paulobilynskyj1 + 311720533, // @katiabacelar 312252804, // @depandresoares @@ -1768,6 +1861,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 343512877, // @MarceloBelinati 345512946, + // @SertanejoUFC + 346580194, // @edusantosdf 348089005, // @leilafonsecapb @@ -1794,6 +1889,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 360649695, // @milkleileite 365210575, + // @Thiagofreitaz + 365552309, // @andresalineiro 367519089, // @Ana_claudiapb @@ -1882,8 +1979,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 487622592, // @LPescinelli 492334281, + // @brunosouzasc + 494268633, // @brauliolaranovo 505339844, + // @caiocordeirotv + 521514005, // @glauberbastos_ 527053144, // @MoisesSantosAc @@ -1904,6 +2005,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 583377940, // @HelioWirbiski 589375704, + // @adautosoutoms + 596446640, // @AfonsoFlorence 599427558, // @CiceroSimplicio @@ -1914,12 +2017,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 608730390, // @MarioEsteves2 610309440, + // @pastorhenriquev + 617468983, // @MateusWesp 618506366, // @elmanooficial 626602522, // @fabiogov55 630758039, + // @rosemodestoms + 632251844, // @helenaduailibe_ 632737286, // @andrewleal2 @@ -1992,6 +2099,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1316758495, // @zecarlospt 1325494376, + // @boettchermarcos + 1339028060, // @CARLOSVALADARE7 1356677952, // @D_GoretePereira @@ -2012,6 +2121,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1532130170, // @Luiza_RibeiroG 1544201047, + // @pedrocampospe + 1546304018, // @natthpaccola 1570635661, // @NegrahLima @@ -2028,6 +2139,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1651852124, // @carlosedrsantos 1685646356, + // @alcidesfernm + 1687046911, // @Brandaveneno 1710385291, // @chaficlays @@ -2064,6 +2177,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2216639570, // @DepFederalMoses 2217650233, + // @andredopradosp + 2250697531, // @ZeniteRosa 2289590857, // @BahMatteuss @@ -2082,6 +2197,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2353403137, // @capitaoassis10 2359974485, + // @marciolabre + 2401917402, // @sheikhrodrigo 2429206439, // @OmarAzizAm_ @@ -2114,6 +2231,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2580412784, // @anadogasoficial 2583179096, + // @depheliolopes + 2585621855, // @profsoniameire 2604536294, // @depjorgesolla @@ -2228,6 +2347,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3092560931, // @rosangelawm 3096479489, + // @ericgustavodf + 3111776739, // @TeonilioBarba 3119378914, // @otavio_camp @@ -2244,6 +2365,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3167874665, // @meire_cruvinel 3205786257, + // @andre_fufuca + 3254270357, // @Marcio_Honaiser 3294107902, // @deputadopriante @@ -2256,12 +2379,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 3335733075, // @drjeanfreire 3342622547, + // @castrothulio + 3354383374, // @PbnConcursos 3357932231, // @brasleiroluc 3366080079, // @moisesbrazpt 3373574517, + // @caformigoni + 3385089700, // @kleybe_morais 3512854216, // @carmelonetobr @@ -2292,6 +2419,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 4204580127, // @_AlessandroSE 4250596815, + // @IrataAbreu + 4309564575, // @DepLucianoMDB 4350613047, // @josuelsantosbch @@ -2302,6 +2431,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 4648157621, // @FofaBorges 4775264669, + // @andrebrandao_rj + 4820124735, + // @SandroFilhoBA + 4836857632, // @ederborgesbr 4871172143, // @emersonosasco @@ -2340,6 +2473,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 724003686863740933, // @ViniciusFerroC 726243481669242880, + // @emersonjarude + 726777982115799040, // @lpbragancabr 728281672731471873, // @brisabracchi13 @@ -2406,6 +2541,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 793261830072311812, // @RafaelMarcal33 793424665431666688, + // @DaniloBalas + 797527121879068672, // @iyagomedeiros 799344884197064704, // @SocorroLac @@ -2438,10 +2575,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 830144764360192002, // @ranallipf 832322309436403712, + // @OseiasVarao + 832603702737379329, // @depeniotatto 834065907743854592, // @peumendonca23 839127759955841024, + // @DaSilvaBelford + 841440085035896834, // @GeneralGirao 841700087143288832, // @viniciusaithsp @@ -2476,6 +2617,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 885883830489554944, // @eucricielle 887083474104049664, + // @paulosafederal + 889464426222563328, // @brunopedralva 889497158491271169, // @jonesmanoel_PE @@ -2508,6 +2651,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 917373635605786624, // @cabojunioamaral 917727546615193601, + // @NoemiaTiago + 918524953561100288, // @wanderley_vieir 921969720546480128, // @GuajajaraSonia @@ -2542,6 +2687,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 958833209961254912, // @DepJorgeEverton 959184931552464896, + // @ClaudioLCaivano + 962035539133177858, // @Fbgg40 963928182368980993, // @RafaelLustoza_ @@ -2624,10 +2771,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1000525959257378816, // @Izalourenca 1000724258740473856, + // @tabataamaralsp + 1001251931812220928, // @NatanSperafico 1001476076466593792, // @zuccors 1002182052341534720, + // @romulobuldrini + 1002646368496865282, // @paparicobacchi 1003614554394415104, // @jmtavaresz @@ -2694,6 +2845,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1032017342937661441, // @vereadornetuno 1032654017418211330, + // @SalvinoOliveir1 + 1034050960262418433, // @victorhugoforte 1034054725107363840, // @RafagninLuciana @@ -2718,6 +2871,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1040704983358955526, // @RenanSantosMBL 1042601099566436352, + // @Daysehansa1 + 1044658383922638849, // @CavalarEmanuel 1047079027696177152, // @gutopfonseca @@ -2742,14 +2897,20 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1055640991200411648, // @cleitinhotmj 1057231251743170562, + // @gilson__marques + 1057602315664924673, // @delegadasheila 1058010509256126464, // @mauricio_lindol 1058341993372360705, + // @evandro_stacruz + 1059551999258255361, // @veronicalima_ve 1059554967600685058, // @pinheirinhomg 1060134845043666945, + // @RuthVenceremos + 1061207070412886016, // @giordanmes 1062407588724252673, // @pluviapt @@ -2770,6 +2931,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1069259896892329984, // @DelHelioBressan 1070740683386957825, + // @wilkerleaods + 1073662759009771520, // @marxbeltrao 1074651587140902913, // @capalbertoneto @@ -2802,6 +2965,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1084272914202091520, // @majorfabianadep 1084712593292443648, + // @waltinhofoguete + 1084870539666247682, // @wilkerbarretoam 1085537170276958208, // @ArthurLira_ @@ -2866,6 +3031,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1098328459825287187, // @helio_ver 1098911814824411137, + // @DelegadoFreitas + 1098981858321276929, // @benesleocadiorn 1099003656064643073, // @BlogdoSavio @@ -2940,6 +3107,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1123714030391132161, // @emanuelvenzo17 1124016358893785089, + // @AltairMoraessp + 1124045706883473408, // @eunatashapoa 1124140833547141121, // @GilbertoCattan1 @@ -2958,6 +3127,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1130259444326109184, // @fabiosilvadep 1130509226281975814, + // @GersonSPires + 1130572020901699585, // @gilmarpetrolina 1131947049690255360, // @brunolessarj @@ -3036,6 +3207,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1179012777366736899, // @delegadopalumbo 1179437585275465729, + // @by_dzn + 1184622911174389762, // @rickazzevedo 1184969015023800320, // @PauloSussumu @@ -3048,6 +3221,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1188827192999972865, // @BrunoZambelli3 1191788412833075201, + // @rebecaromerobr + 1192983475663638529, // @ProtetorAle 1193938478456868865, // @thiagomedinamd @@ -3126,8 +3301,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1227572253522526208, // @Alfredogaspar_ 1228008951813492739, + // @izadoradias29 + 1231267069762691072, // @BabaTupinamba 1232082766071767045, + // @felipegomespsol + 1232252941681221632, // @amandasalecosta 1232788071218843648, // @queciareismbl @@ -3178,12 +3357,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1244783574676627460, // @docporto 1244954874049114113, + // @allinyserrao + 1245187621007089666, // @vitimporto 1246125973755609088, // @wmf_oficial 1246641106391052288, // @GallinatiRaquel 1246945909881044994, + // @RodrigoLivra + 1247224255311499264, // @drluizovando 1247287375803355136, // @robertlemosss @@ -3250,12 +3433,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1262645801936879617, // @nise_dra 1262944997118181385, + // @mikaelgringoPE + 1265422039260762113, // @hana_ghassan 1266449427985829888, // @professorarita_ 1266477339413749761, // @mariadoscamelos 1267127677468753920, + // @Junacamara + 1267549843586768896, // @aanaelisast 1267822695195930630, // @ChirleyPankara @@ -3368,6 +3555,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1310977778574032900, // @PabloSilvaLira 1311620519935045632, + // @OdairTramontin + 1311753908549750784, // @gedalvaumbauba 1312771507601498112, // @cassymonteiro @@ -3386,12 +3575,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1322213572953493505, // @PamelaGiedre 1325859971905572869, + // @bfeministapsol + 1329213002051149824, // @paulomelo_sa 1330155785951866881, // @cabomeireles 1330393555290951681, // @patisborges 1330690808090124289, + // @_PastorDiego + 1334311453592064000, // @bocalomoficial 1335685798746845190, // @RoseanaSarneyM @@ -3458,12 +3651,18 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1362442787988381696, // @boulosnat 1366796662614749187, + // @NeidinhaSurui + 1368009955912130561, // @_mapeoficial 1372519763818217477, + // @NilmaSanches + 1373064659817926661, // @rodrigoestacho 1374044251475091457, // @PiauienseO 1374436642375659530, + // @VitorioJrMG + 1374502893533884421, // @FelipeAlecrimPE 1375498458132537348, // @rodolfoms @@ -3474,6 +3673,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1379025327314329608, // @DrLeviMelo1 1379424037068234752, + // @fabiosilveirarn + 1379583456812924935, // @marleipr 1380366998488645636, // @depprofcleiton @@ -3486,6 +3687,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1385208512955957250, // @SylvioMenicucci 1385213373332213778, + // @PCO29antonio + 1387843668053200897, // @depmarcionunes 1390475281316646917, // @davibrandaobac1 @@ -3498,6 +3701,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1395784254269841408, // @DelegadoMarcus 1396055754793230341, + // @renataseneofc + 1399347605776314368, // @jusoaresft 1400452986178981891, // @eduacostario @@ -3516,6 +3721,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1409194659872653317, // @sauloportolivei 1409385758838906882, + // @jrfreitasofc_ + 1409935535443951616, // @juniorferrari55 1409990621536919555, // @rafaelsaraivasp @@ -3548,6 +3755,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1434501105241792512, // @NFgoes 1435915606541410305, + // @0xcarlosh + 1437029285387243526, // @PeKelmon 1437437148769226757, // @ScalcoDarlan @@ -3566,6 +3775,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1445155253880594432, // @renatmirandarj 1446451214292553733, + // @DayaneCruzPE + 1446618743350730754, // @DanielaadvAP 1446856314215337989, // @joelrodriguespi @@ -3590,6 +3801,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1457433413028356098, // @fala_mafia 1458089891795976202, + // @sargentoportuga + 1458495272716222465, // @josecam01577970 1459681826700673030, // @schumarker7 @@ -3604,6 +3817,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1463574028468330497, // @padovanidep 1464329410115514369, + // @michelkourimbl + 1464589158035468306, // @marinadomst 1465395158103597065, // @amarianalescano @@ -3612,6 +3827,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1468757430632865793, // @Maubmarcon 1469007240279597067, + // @orenatomachado + 1471286780737503232, // @joaquimroriznet 1472329069027119107, // @LuizinhoMinas @@ -3724,10 +3941,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1516033327941181447, // @DanielleDVale 1516121393028644869, + // @thiagorodripa + 1516560691410481156, // @andrebuenoofc 1517514873105797120, // @ivanilsonrn 1517576201338040321, + // @leandrobasson + 1518652449371922433, + // @drpaulomedinamg + 1518728271176871936, // @EDILSONLIMA70 1518759485694779392, // @EdRaposo_ @@ -3780,6 +4003,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1530904551582220288, // @JadyelAlencarr 1531271393219837956, + // @Jaimenunes_ap + 1531655976008531968, // @ReginaluciaSi15 1531933026208436231, // @eudonaneuma @@ -3836,6 +4061,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1550186795333296128, // @julio_kuller 1551652560951562240, + // @DepAndreAbdon + 1552328294665781249, // @Gi_MonteiroRJ 1553042102908502021, // @MissiasDias @@ -3852,6 +4079,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1556676552674394114, // @KleberRosa50 1556777705387032577, + // @josyjacob_ + 1556827882516893700, // @VitormoreiraCG 1557086339740405761, // @ScaranteRenato @@ -3890,10 +4119,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1575549757900132379, // @14luizfranca 1575672365094440960, + // @PauloEchebarria + 1576793388845760512, // @ProfAlexandreS 1577470375272878083, // @Arnaldodeputado 1578127734362046489, + // @Mariada67416358 + 1578375608878391296, // @TamirisPeixoto1 1578450870127263757, // @CarlosValdevin6 @@ -3914,6 +4147,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1584299141550702593, // @andreadantasc 1584527735686397952, + // @tcoronel_maia + 1586016154262315009, // @rafandradembl 1586192715137667072, // @CombatPatriota @@ -3932,6 +4167,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1589336373101723649, // @JUNIORCESARLEI9 1589412695664742401, + // @LuanNun43052492 + 1589645603616817154, // @pitypaguiar 1589741307630690304, // @Marisaloboreal @@ -3974,6 +4211,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1605547416366841856, // @MarinaCallega13 1609042267095941123, + // @gutembergfrios + 1610745043358187520, // @WickRyanAM 1612178912947183620, // @SchiavoMaurilio @@ -3988,6 +4227,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1612589448256016385, // @gersonzocchi1 1612930692932837377, + // @oGabrielAzevedo + 1612935067487145985, // @WaldenorPereira 1613183981230366728, // @DiegoQuaqua @@ -4024,12 +4265,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1634553908402987010, // @Dep_CoronelNeil 1635697153979891712, + // @LucasAl29111516 + 1637240612620587013, // @AdrianaLeal2507 1637963109573828608, // @adrianalmeidapt 1638181881127612416, // @leticiaaguiarsp 1638274776413134849, + // @Rogerio_Ulysses + 1641171155540123655, // @edercostarj 1642933716279328768, // @giuargolo @@ -4042,6 +4287,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1648046380106104833, // @Alannleal_ 1651297457374887936, + // @felixsantosrn + 1653188395324014593, // @CamilaGodoiSP 1653466465855471617, // @MatiasSamuka @@ -4094,6 +4341,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1708844117914972162, // @willianrochapr 1709986969277624320, + // @osenhorsaldanha + 1710320865517162496, // @francisco_arten 1710430261932908545, // @hebertcsgyn @@ -4176,6 +4425,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1777608369236271105, // @igorrayanrn 1777621309335183360, + // @Dhe_paiv + 1777901547629764608, // @juhliasantost 1778077250908274688, // @OficialVanucci @@ -4186,6 +4437,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1779913435825795072, // @BiaCoimbraSP 1780537023859761153, + // @PedroNassifrj + 1780805773603258368, + // @paulacoutinho65 + 1781108419853701120, // @CamillaGonda 1781485235047174144, // @brunnomattospt @@ -4204,6 +4459,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1789440263187771392, // @GuimaraesAlpha 1790692252840243201, + // @EloyNhandewa + 1791120964731744256, // @coronellrosses 1793428861348167680, // @akarinaclaro @@ -4224,6 +4481,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1801749578938564608, // @leorondonn13 1803570674914426880, + // @igor_ooliveira1 + 1809003112658767872, // @gabi_bvnt 1809421300433317892, // @ustramarcelo22 @@ -4298,6 +4557,10 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1859357204278542336, // @LusFelipeV68715 1861380355875373056, + // @lucylulv + 1863493161240133632, + // @_AnaHering + 1864125622928187392, // @denistaveiradn 1865917852424798208, // @joaopaulo_tprs @@ -4366,6 +4629,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1904187270342574082, // @profeVinicius 1906160472354873344, + // @ErikaGuerrieri + 1906823280281174016, // @rodrigospada_ 1908175828417949697, // @MarinhoGui65411 @@ -4384,6 +4649,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1920469901136764928, // @ericafmissao26 1922067675074727936, + // @joselaviola + 1922378485412331522, // @daniele_carva 1925558068206518272, // @MuriloM50730958 @@ -4398,10 +4665,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1932924340917710848, // @Camargo1Amanda 1937005496398962688, + // @JulianaMontrios + 1937336667356205056, // @prof_elson_sc 1937980838299504644, // @marcioalvinosp 1940410346411597824, + // @ArturdeFariasV + 1941922310568443904, // @beto_vaz_ 1944257271942287360, // @jotabrandaoam @@ -4412,6 +4683,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1953095258784579586, // @IgorMarquesBRA 1957556152117653505, + // @lucasvenditeMS + 1958412111752699904, // @NoronhaPro7153 1958622002719182848, // @RubensAngiolett @@ -4434,8 +4707,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1966489240566444033, // @edilsondamiaorr 1970177461070528512, + // @joaoadolfo_14 + 1970451742430035968, // @Drfilipecm 1972344363654037504, + // @delsmontanari + 1973824020907700224, // @romulobraz_13 1975026318975840257, // @DrManuelMarcos @@ -4452,6 +4729,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1980794408141258752, // @drjuliostobbe 1981365553173299200, + // @Kiko_Caputo + 1982938414002647045, // @mazzei_fel47870 1983280886067195904, // @marciaabrahaodf @@ -4480,14 +4759,22 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 1997124812662362113, // @oleandroissa 1998066982999212032, + // @Euberlucas14 + 1998418256953233408, // @profakellsilva 1998797409712009216, + // @joaopiresrj + 2000205916009074690, // @nelsongrasselli 2002006532184281088, + // @rafoliver_pa + 2005281365072449536, // @moisesbarboza 2006965910309896192, // @alvarenga35289 2007441310102523904, + // @coronelbusnello + 2009066265156202496, // @RogerioChimi 2010017755194589185, // @AntoniadeJessp @@ -4496,6 +4783,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2010532868271816705, // @profraydf 2010751427623497728, + // @PolianadoPsol + 2010869406650245120, // @edinhosouzaaa 2011825328793296896, // @Efreu_Quintana @@ -4532,8 +4821,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2021639549625991168, // @henriqueduraesd 2021718062261506049, + // @Jordambritosc + 2022354011035222016, // @sophiafechinece 2022429125399478273, + // @mateusquadrosms + 2024092564014669824, + // @tulioteconta + 2024178238990536704, // @NetoFeitos68916 2026461466728292352, // @catarinanevespb @@ -4550,10 +4845,14 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2029985728059559938, // @helen_vitaRJ 2033587304795963392, + // @jeaninepresente + 2034019355629846528, // @CidaCarvalhoSP 2034086382705274880, // @EdneyBatalha 2034242236678848513, + // @WickRyaAM + 2034644862679547904, // @Fabio_x86 2036151026986696704, // @MarcaoVivacqua @@ -4580,6 +4879,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2041763427488608256, // @Grazypasqualeto 2042058510536482816, + // @vivianegomesrn + 2042331395725430784, // @CrisNavarroMIDH 2043108953177899008, // @vanessacfortes @@ -4600,6 +4901,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2047324463612485632, // @JoaoPaulo_2026 2047395841715920897, + // @DeyconSP + 2048379106517987328, // @VasconcellosCel 2048607450593468416, // @PablodoMST @@ -4612,6 +4915,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2049539394701262848, // @bia_pedagoga 2049646600088117249, + // @enzodangelo_mbl + 2049839952624492544, // @oalanmansurrj 2049927296442658816, // @14DanielAguiar @@ -4632,6 +4937,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2054954070083813376, // @mahmoudamer_rs 2054974955603861504, + // @luizgamabsb + 2055147146618302464, // @DelEduardoK 2055437449740873728, // @glaucelima12 @@ -4640,6 +4947,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2056559848750276608, // @celmarcioasouza 2056815030935416832, + // @CidaVillasBoas + 2056824758201708544, // @RodolfoFiorucci 2057456621865881600, // @brenobarcelos_ @@ -4666,6 +4975,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2060549322207375368, // @nicolasravipsol 2061280727711289345, + // @Renatofonsecape + 2061591502136983552, // @DrCrisVeloso 2061889166779019264, // @DaversonMatos @@ -4702,6 +5013,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2069869784972369920, // @Waltinhogo 2070185289100775424, + // @candidotelesdr + 2070216303479066624, // @Deborahzanchi_ 2070650915871240192, // @Delmariacorsato @@ -4712,6 +5025,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2071779601815134208, // @angelaabukce 2071949683736354816, + // @wilsonamigao + 2071988583334862849, // @Drrafujr 2074177828606652416, // @rogeriozabdalla @@ -4772,6 +5087,8 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2080632627887734784, // @thiagocampeloof 2081154349267345408, + // @erigreyce__ + 2081487029443940352, // @DanielSantanamc 2081774169167974415, // @luismartarj @@ -4792,12 +5109,16 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2082567723423342592, // @Clarianabr 2082581908005806081, + // @joaogusmao_pt + 2082766963659403264, // @karimeefayad 2082897339065233409, // @barbararesende0 2082924803145551872, // @guihenriquesc 2082932165466058752, + // @odrlucasoficial + 2083260564294295552, // @JHONNSOM70 2083921524394700800, // @eliethdefatima @@ -4830,8 +5151,12 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2085133776741437441, // @diegojejees 2085201272139886592, + // @profyoshida + 2085344308480049152, // @lucianoleitoa_ 2085736572914159616, + // @Virginiaenfer + 2085739235093446656, // @profandersonfig 2085812265509388288, // @nandovianapsol @@ -4854,10 +5179,20 @@ static BRAZIL_2026_ELECTION_USER_IDS: LazyLock> = LazyLock::new(| 2089024802669334528, // @brunoscheid222 2089064009643220993, + // @RafaGonzagaSP + 2089736144921444352, + // @marciacandidata + 2089837151844212736, // @SimonePimehb 2090082297370312704, + // @RadmanGadiel + 2090252697354067968, // @cabodaciolo33 2091177221310345216, + // @danilosoaresce + 2091948027967676416, + // @IuriMarque60kt + 2092800629278126081, ]) }); @@ -5045,7 +5380,7 @@ mod tests { #[test] fn hardcoded_list_is_non_empty() { assert!(!BRAZIL_2026_ELECTION_USER_IDS.is_empty()); - assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2388); + assert_eq!(BRAZIL_2026_ELECTION_USER_IDS.len(), 2554); } #[test] diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index 7b3825c5..7fd5eb54 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-09-02T16:03:45Z +// mirrored from config feature-switch defaults; last sync 2026-09-03T16:23:24Z use xai_feature_switches::param; param!( @@ -625,6 +625,12 @@ param!( "rust_home_mixer_vm_ranker_cluster_id", "Experiment3" ); +param!( + PhoenixExperimentOverrides, + String, + "rust_home_mixer_phoenix_experiment_overrides", + "" +); param!( VMRankerDppTheta, f64, diff --git a/home-mixer/server.rs b/home-mixer/server.rs index ed51d90e..d8f59d18 100644 --- a/home-mixer/server.rs +++ b/home-mixer/server.rs @@ -387,7 +387,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { - let polling_header = strato_context::is_polling(request.metadata()); + 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 @@ -395,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 = polling_header || proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -414,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)) => { @@ -480,7 +482,7 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { &self, request: Request, ) -> Result, Status> { - let polling_header = strato_context::is_polling(request.metadata()); + let strato_ctx = strato_context::parse(request.metadata()).unwrap_or_default(); let mut b3_info = extract_b3_info(request.metadata()); b3_info.force_sample(); @@ -491,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 = polling_header || proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -511,6 +513,8 @@ impl pb::for_you_feed_service_server::ForYouFeedService for ForYouFeedServer { 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)) => { @@ -544,7 +548,7 @@ impl pb::ranked_following_feed_service_server::RankedFollowingFeedService &self, request: Request, ) -> Result, Status> { - let polling_header = strato_context::is_polling(request.metadata()); + 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 @@ -552,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 = polling_header || proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -572,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)) => { @@ -645,7 +651,7 @@ impl pb::following_feed_service_server::FollowingFeedService for FollowingFeedSe &self, request: Request, ) -> Result, Status> { - let polling_header = strato_context::is_polling(request.metadata()); + 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 @@ -653,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 = polling_header || proto_query.is_polling; + let is_polling = strato_ctx.is_polling || proto_query.is_polling; let ctx = self .query_builder .build( @@ -673,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)) => { diff --git a/home-mixer/util/phoenix_request.rs b/home-mixer/util/phoenix_request.rs index 16779871..222a6ad7 100644 --- a/home-mixer/util/phoenix_request.rs +++ b/home-mixer/util/phoenix_request.rs @@ -1,6 +1,6 @@ use crate::models::candidate::{CandidateHelpers, PostCandidate}; use crate::models::query::ScoredPostsQuery; -use crate::params::RerankerHeadTag; +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; @@ -136,10 +136,23 @@ pub fn build_request_without_sequence_and_candidates( 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], diff --git a/home-mixer/util/strato_context.rs b/home-mixer/util/strato_context.rs index 44494ca8..c7a85b39 100644 --- a/home-mixer/util/strato_context.rs +++ b/home-mixer/util/strato_context.rs @@ -6,16 +6,16 @@ const STRATO_CONTEXT_KEY: &str = "stratocontext"; const STRATO_CONTEXT_BIN_KEY: &str = "stratocontext-bin"; #[derive(Clone, PartialEq, prost::Message)] -struct StratoContext { +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 is_polling(metadata: &MetadataMap) -> bool { - extract_strato_context(metadata).is_some_and(|ctx| ctx.is_polling) -} - -fn extract_strato_context(metadata: &MetadataMap) -> Option { +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()) @@ -44,18 +44,23 @@ mod tests { map } + fn polling(ctx: StratoContext) -> StratoContext { + StratoContext { + is_polling: true, + ..ctx + } + } + #[test] fn polling_true() { - assert!(is_polling(&encode_ascii(&StratoContext { - is_polling: true - }))); + let ctx = parse(&encode_ascii(&polling(StratoContext::default()))).unwrap(); + assert!(ctx.is_polling); } #[test] fn polling_false() { - assert!(!is_polling(&encode_ascii(&StratoContext { - is_polling: false - }))); + let ctx = parse(&encode_ascii(&StratoContext::default())).unwrap(); + assert!(!ctx.is_polling); } #[test] @@ -63,13 +68,25 @@ mod tests { let mut map = MetadataMap::new(); map.insert_bin( STRATO_CONTEXT_BIN_KEY, - MetadataValue::from_bytes(&StratoContext { is_polling: true }.encode_to_vec()), + MetadataValue::from_bytes(&polling(StratoContext::default()).encode_to_vec()), ); - assert!(is_polling(&map)); + assert!(parse(&map).unwrap().is_polling); } #[test] fn missing_context() { - assert!(!is_polling(&MetadataMap::new())); + 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/phoenix/crates/common/xai-recsys/src/util.rs b/phoenix/crates/common/xai-recsys/src/util.rs index ceb672e7..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,6 +29,75 @@ lazy_static! { .unwrap(); } +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"]) @@ -1546,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, diff --git a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml index 91ecc659..a2a993ad 100644 --- a/phoenix/crates/serving/xai-recsys-engine/Cargo.toml +++ b/phoenix/crates/serving/xai-recsys-engine/Cargo.toml @@ -62,7 +62,6 @@ 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 } 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/checkpoint_proxy.rs b/phoenix/crates/serving/xai-recsys-engine/src/checkpoint_proxy.rs index 03ba070a..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,14 +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).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; @@ -332,11 +399,10 @@ impl CheckpointProxy { download_channel_files( &staging_dir, &grpc_prefix, - channel, + pool, &ch_entries, ch_idx, verify, - concurrency, ) .await }); @@ -481,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; @@ -501,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 { @@ -546,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(); @@ -578,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, @@ -588,7 +656,7 @@ async fn download_channel_files( } else { 0.0 }, - concurrency, + n_conns, ); Ok(()) @@ -645,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( @@ -683,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 {}-{}: {:?}", @@ -694,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/emb_table.rs b/phoenix/crates/serving/xai-recsys-engine/src/emb_table.rs index 5a561b50..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>>, enqueue_timeout_ms: u64, mm_embeddings_client: Option, - #[allow(dead_code)] - sid_client: Option>, admission: Arc, #[allow(dead_code)] prefetch_mm_query_config: Option>, @@ -2734,8 +2732,6 @@ struct RecsysRetrievalPredictorImpl { reload_directive: Arc>>, enqueue_timeout_ms: u64, mm_embeddings_client: Option, - #[allow(dead_code)] - sid_client: Option>, admission: Arc, prefetch_mm_query_config: Option>, } @@ -3097,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], @@ -3162,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, @@ -3311,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() { @@ -3350,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, }; @@ -3689,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/sid_client.rs b/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs deleted file mode 100644 index be094e00..00000000 --- a/phoenix/crates/serving/xai-recsys-engine/src/sid_client.rs +++ /dev/null @@ -1,189 +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, -} - -pub fn fill_semantic_ids( - sids: &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; - } - } - } -} - -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, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn semantic_ids_are_shifted_and_missing_values_remain_padding() { - let sids = HashMap::from([(11, vec![0, 7, -1]), (22, vec![3, 4, 5])]); - let mut dst = vec![0u16; 3 * 3]; - - fill_semantic_ids(&sids, &[11, 0, 22], &mut dst, 3); - - assert_eq!(dst, vec![1, 8, 0, 0, 0, 0, 4, 5, 6]); - } -} 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 1c37809c..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] = ..., @@ -279,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] = ..., diff --git a/phoenix/xrex/cutedsl/ranker_attention_fa4.py b/phoenix/xrex/cutedsl/ranker_attention_fa4.py index 88bffaee..47d0fe68 100644 --- a/phoenix/xrex/cutedsl/ranker_attention_fa4.py +++ b/phoenix/xrex/cutedsl/ranker_attention_fa4.py @@ -10,69 +10,95 @@ _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 @@ -94,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: @@ -111,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, ) @@ -143,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, @@ -159,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, @@ -219,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, @@ -236,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, @@ -269,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), ) @@ -351,20 +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) - 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: @@ -378,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/inference/launch_inference.py b/phoenix/xrex/inference/launch_inference.py index d0cc1dc5..bd109b5e 100644 --- a/phoenix/xrex/inference/launch_inference.py +++ b/phoenix/xrex/inference/launch_inference.py @@ -384,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: @@ -649,9 +647,9 @@ 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", @@ -804,15 +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=( - "Optional leftover SID lookup endpoint. History SIDs are parsed " - "from the request; this is not required for use_post_sid=True." - ), - ) parser.add_argument( "--jax_compilation_cache_dir", type=str, diff --git a/phoenix/xrex/inference/model_runner.py b/phoenix/xrex/inference/model_runner.py index e62ce54c..ed432442 100644 --- a/phoenix/xrex/inference/model_runner.py +++ b/phoenix/xrex/inference/model_runner.py @@ -501,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 @@ -515,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 @@ -4747,24 +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 else 0 - 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: - logger.info( - "Parsing history SIDs from the request (sid_num_levels=%d); no sid_endpoint", - sid_num_levels, - ) + 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, @@ -4781,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, diff --git a/phoenix/xrex/inference/sid_retrieval_runner.py b/phoenix/xrex/inference/sid_retrieval_runner.py index c20149dd..095ea3d5 100644 --- a/phoenix/xrex/inference/sid_retrieval_runner.py +++ b/phoenix/xrex/inference/sid_retrieval_runner.py @@ -377,23 +377,7 @@ def create_server( assert isinstance(self.model_config, RecsysSIDRetrievalConfig) hash_keys = self.dataset.hash_table.hash_keys - sid_client = None sid_num_levels = self.model_config.sid_num_levels if self.model_config.use_post_sid else 0 - if 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 self.model_config.use_post_sid: - logger.info( - "Parsing history SIDs from the request (sid_num_levels=%d); no sid_endpoint", - sid_num_levels, - ) return xai_recsys_engine.RecsysRetrievalPredictorServer( self.grpc_port, @@ -410,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, diff --git a/phoenix/xrex/models/recsys_attention.py b/phoenix/xrex/models/recsys_attention.py index 77c2c06f..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, () 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/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 be32bb26..55c5466e 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -1,12 +1,6 @@ -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_DUAL_CALL_HARNESS_ENABLED: &str = "VF_DUAL_CALL_HARNESS_ENABLED"; -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_CACHE_WARM_SAMPLE_PCT: &str = "VF_CACHE_WARM_SAMPLE_PCT"; +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"; @@ -46,35 +40,12 @@ pub fn dual_call_harness_enabled() -> bool { parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) } -pub fn fallback_cache_serve_stale_enabled() -> bool { - parse_env_flag( - std::env::var(ENV_FALLBACK_CACHE_SERVE_STALE_ENABLED) - .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 fn fallback_cache_enabled() -> bool { + parse_env_flag(std::env::var(ENV_FALLBACK_CACHE_ENABLED).ok().as_deref()) } -pub fn cache_warm_sample_pct() -> u8 { - parse_sample_pct(std::env::var(ENV_CACHE_WARM_SAMPLE_PCT).ok().as_deref()) - .unwrap_or_else(|error| panic!("{ENV_CACHE_WARM_SAMPLE_PCT}: {error}")) -} - -fn parse_sample_pct(value: Option<&str>) -> Result { - let Some(value) = value.map(str::trim).filter(|v| !v.is_empty()) else { - return Ok(0); - }; - match value.parse::() { - Ok(pct) if pct <= 100 => Ok(pct), - _ => Err(format!("expected an integer 0-100, got {value:?}")), - } +pub 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 { @@ -86,91 +57,9 @@ 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, parse_sample_pct, resolve_gizmoduck_client_id, - resolve_twemcache_client_name, - }; - - #[test] - fn parses_sample_pct_range() { - assert_eq!(parse_sample_pct(None), Ok(0)); - assert_eq!(parse_sample_pct(Some("")), Ok(0)); - assert_eq!(parse_sample_pct(Some("100")), Ok(100)); - } - - #[test] - fn rejects_invalid_sample_pct() { - for value in ["101", "on"] { - assert!(parse_sample_pct(Some(value)).is_err(), "{value}"); - } - } + use super::{parse_env_flag, resolve_gizmoduck_client_id, resolve_twemcache_client_name}; #[test] fn client_ids_default_to_historical_values_and_overrides_win() { diff --git a/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs index 8626a983..f37881bf 100644 --- a/visibility-filtering/dark_traffic_setup.rs +++ b/visibility-filtering/dark_traffic_setup.rs @@ -181,6 +181,7 @@ pub fn resolve_layer() -> DarkLayer { let domain = staging_tls_domain(&dc); info!(domain, "dark_traffic: enabled"); + #[expect(clippy::expect_used, reason = "startup fail-fast: TLS is required")] let factory = XdsChannelFactory::new( TlsMode::mtls_from_env() .expect("S2S TLS config required") diff --git a/visibility-filtering/filter.rs b/visibility-filtering/filter.rs index 032b56a1..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::{RuleEngine, SafetyLevel, Verdict}; use std::collections::HashMap; -use tracing::{debug, info}; use xai_visibility_filtering_proto as vf_pb; pub struct FilterRequest { @@ -19,15 +18,8 @@ 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 { @@ -58,14 +50,6 @@ 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| { @@ -85,14 +69,6 @@ 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, @@ -108,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)] 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; @@ -179,7 +145,7 @@ pub(crate) mod test_support { gizmoduck: Arc, ) -> FilterTweets { let tes: Arc = Arc::new(MockTESClient::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( @@ -187,14 +153,8 @@ pub(crate) mod test_support { )))); FilterTweets::new( - HydrationPipeline::new( - tes, - gizmoduck, - socialgraph, - labels, - crate::hydration::FallbackCacheMode::Disabled, - ), - RuleEngine::new(), + HydrationPipeline::new(tes, gizmoduck, socialgraph, labels, None), + RuleEngine::for_tests(), ) } } @@ -203,6 +163,7 @@ pub(crate) mod test_support { 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 { @@ -250,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 79aa575a..fe653ff6 100644 --- a/visibility-filtering/filter_tweets.rs +++ b/visibility-filtering/filter_tweets.rs @@ -4,9 +4,7 @@ use crate::reference_compare::{ReferenceCompareHarness, TweetVerdict}; use crate::rules::metrics::{self as ft_metrics, RequestMetricsGuard}; use crate::rules::SafetyLevel; use std::sync::Arc; -use std::time::Instant; use tonic::{Request, Response, Status}; -use tracing::info; use xai_visibility_filtering_proto as vf_pb; pub struct FilterTweetsEndpoint { @@ -30,7 +28,6 @@ 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); @@ -43,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() @@ -98,14 +87,6 @@ impl FilterTweetsEndpoint { .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 })) } @@ -191,51 +172,22 @@ 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()); - } - - #[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)); - } - - #[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 f9630b40..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, @@ -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() } @@ -189,14 +161,10 @@ mod tests { ) .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] @@ -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,16 +210,12 @@ 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] 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 1e716aa3..7f905586 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; @@ -135,7 +134,7 @@ impl HydrationPipeline { gizmoduck_client: Arc, socialgraph_client: Arc, safety_label_source: Arc, - fallback_cache_mode: FallbackCacheMode, + fallback_cache: Option>, ) -> Self { Self { viewer_hydrator: ViewerHydrator { @@ -146,7 +145,7 @@ impl HydrationPipeline { }, gizmoduck_author_hydrator: GizmoduckAuthorHydrator::new( GizmoduckLookup::new(gizmoduck_client), - fallback_cache_mode, + fallback_cache, ), socialgraph_hydrator: SocialgraphHydrator { sg_client: socialgraph_client.clone(), @@ -173,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) }; @@ -296,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) }; 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 cf873fd5..fb5f21c5 100644 --- a/visibility-filtering/hydration/tes_hydrator.rs +++ b/visibility-filtering/hydration/tes_hydrator.rs @@ -1,8 +1,7 @@ use crate::hydration::batch::TweetHydrationBatch; 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; @@ -24,25 +23,11 @@ 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 async fn fetch_pure_core( &self, @@ -83,7 +68,6 @@ impl TesHydrator { community, nsfw_user, nsfw_admin, - has_takedown, takedown_reasons, edit_control, media_entities, @@ -120,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", @@ -159,7 +135,6 @@ impl TesHydrator { 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_entities.map_keys(TweetId).map(media_feature), @@ -202,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), @@ -218,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, @@ -480,7 +451,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); } } diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index 4b967ea1..4e3e44b0 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -1,3 +1,27 @@ +#![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 mod clients; pub mod config; pub mod dark_traffic_setup; diff --git a/visibility-filtering/main.rs b/visibility-filtering/main.rs index 4d6c4c99..19ab936c 100644 --- a/visibility-filtering/main.rs +++ b/visibility-filtering/main.rs @@ -1,3 +1,15 @@ +#![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; diff --git a/visibility-filtering/models/mod.rs b/visibility-filtering/models/mod.rs index 6fbfc0bf..237aaacd 100644 --- a/visibility-filtering/models/mod.rs +++ b/visibility-filtering/models/mod.rs @@ -8,8 +8,8 @@ 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 safety_labels::{SafetyLabelMap, SafetyLabelType}; +pub use tweet::{CoreFeature, MediaFeature, NsfwFeature, TweetFeatures}; pub use viewer::{Viewer, ViewerAge, ViewerFeatures, ADULT_AGE_YEARS}; use std::collections::HashMap; diff --git a/visibility-filtering/models/safety_labels.rs b/visibility-filtering/models/safety_labels.rs index cb7e0b7d..2ef75f06 100644 --- a/visibility-filtering/models/safety_labels.rs +++ b/visibility-filtering/models/safety_labels.rs @@ -1,34 +1,28 @@ -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, - } + 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..f6f6a922 100644 --- a/visibility-filtering/models/tweet.rs +++ b/visibility-filtering/models/tweet.rs @@ -2,7 +2,6 @@ pub struct CoreFeature { pub text: String, pub source_tweet_id: Option, - pub created_at_secs: Option, } #[derive(Clone, Debug, Default)] @@ -13,12 +12,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 +22,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/params.rs b/visibility-filtering/params.rs index af4dcf3d..453929de 100644 --- a/visibility-filtering/params.rs +++ b/visibility-filtering/params.rs @@ -26,7 +26,7 @@ pub struct NsfwGatingCountries { } impl NsfwGatingCountries { - pub fn new() -> Self { + pub fn starting_at_default() -> Self { Self { countries: ArcSwap::from_pointee(default_nsfw_gating_countries()), } @@ -36,6 +36,7 @@ impl NsfwGatingCountries { 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)); @@ -63,12 +64,6 @@ impl NsfwGatingCountries { } } -impl Default for NsfwGatingCountries { - fn default() -> Self { - Self::new() - } -} - fn lowercased_codes(values: &[Value]) -> Option> { values .iter() @@ -149,7 +144,7 @@ mod tests { #[test] fn refresh_reads_key_and_fails_open() { - let cache = NsfwGatingCountries::new(); + let cache = NsfwGatingCountries::starting_at_default(); assert!(cache.contains("de")); assert!(!cache.contains("xx")); @@ -173,7 +168,7 @@ rust_vf: #[test] fn malformed_value_falls_back_whole_not_partial() { - let cache = NsfwGatingCountries::new(); + let cache = NsfwGatingCountries::starting_at_default(); cache.refresh_from(&engine( r#" rust_vf: diff --git a/visibility-filtering/reference_compare.rs b/visibility-filtering/reference_compare.rs index 86040c0d..ba878a47 100644 --- a/visibility-filtering/reference_compare.rs +++ b/visibility-filtering/reference_compare.rs @@ -200,6 +200,7 @@ fn group_diffs(diffs: &[Diff]) -> Vec> { }); groups.len() - 1 }); + #[expect(clippy::indexing_slicing, reason = "at indexes a group already pushed")] groups[at].tweet_ids.push(diff.tweet_id); } groups @@ -253,19 +254,21 @@ pub(crate) fn chunk_lines( .to_string() .len(); let budget = LINE_BUDGET_BYTES.saturating_sub(header_len); - let mut pages: Vec> = vec![Vec::new()]; + 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 && pages.last().is_some_and(|page| !page.is_empty()) { - pages.push(Vec::new()); + if used + cost > budget && !current.is_empty() { + pages.push(std::mem::take(&mut current)); used = 0; } used += cost; - pages.last_mut().expect("pages is never empty").push(slice); + current.push(slice); } } + pages.push(current); let total = pages.len(); pages .into_iter() @@ -392,7 +395,10 @@ impl ReferenceCompareHarness { harness.emit(safety_level, &counts); if !diffs.is_empty() { for line in chunk_lines(&context, &batch_id(), &diffs) { - println!("{line}"); + #[expect(clippy::print_stdout, reason = "stdout is the diff sink")] + { + println!("{line}"); + } } } }); diff --git a/visibility-filtering/rules/author_rules.rs b/visibility-filtering/rules/author_rules.rs index 51ff97b4..1ebf83a6 100644 --- a/visibility-filtering/rules/author_rules.rs +++ b/visibility-filtering/rules/author_rules.rs @@ -169,34 +169,10 @@ pub(super) const SOCIALGRAPH_DROPS: &[RuleSpec] = &[ #[cfg(test)] mod tests { use super::*; - use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, VfAction, ViewerAuthorRelationship, ViewerFeatures, + use crate::models::{AuthorFeatures, ViewerAuthorRelationship}; + use crate::rules::fixtures::{ + assert_allows, assert_drops, author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID, }; - use crate::rules::fixtures::{author_viewer, candidate, logged_out_viewer, viewer, VIEWER_ID}; - use crate::rules::test_context; - - 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() - ); - } - - 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() - ); - } fn author_flag_features(name: &str) -> AuthorFeatures { let mut features = AuthorFeatures::default(); diff --git a/visibility-filtering/rules/context.rs b/visibility-filtering/rules/context.rs index 3900bae1..5b9b939a 100644 --- a/visibility-filtering/rules/context.rs +++ b/visibility-filtering/rules/context.rs @@ -1,11 +1,9 @@ use crate::models::{HydratedTweetCandidate, SafetyLabelType, ViewerFeatures}; use crate::params::NsfwGatingCountries; -use crate::rules::registry::SafetyLevel; use xai_core_entities::entities::TakedownReason; use xai_x_thrift::user_labels::LabelValue; pub struct RuleContext<'a> { - safety_level: SafetyLevel, viewer: &'a ViewerFeatures, candidate: &'a HydratedTweetCandidate, nsfw_gating_countries: &'a NsfwGatingCountries, @@ -13,24 +11,17 @@ pub struct RuleContext<'a> { impl<'a> RuleContext<'a> { pub(super) fn new( - safety_level: SafetyLevel, viewer: &'a ViewerFeatures, candidate: &'a HydratedTweetCandidate, nsfw_gating_countries: &'a NsfwGatingCountries, ) -> Self { Self { - safety_level, viewer, candidate, nsfw_gating_countries, } } - #[inline] - pub fn safety_level(&self) -> SafetyLevel { - self.safety_level - } - #[inline] pub fn viewer(&self) -> ViewerPredicates<'_> { ViewerPredicates { ctx: self } @@ -268,17 +259,18 @@ impl TakedownPredicates<'_> { #[inline] fn in_viewer_country(&self, extractor: fn(&TakedownReason) -> Option<&str>) -> bool { - let Some(viewer_country) = &self.ctx.viewer.country_code else { - return false; - }; + let viewer_country = self.ctx.viewer.country_code.as_deref(); self.ctx .candidate .tweet_features - .takedown - .reasons + .takedown_reasons .iter() .filter_map(extractor) - .any(|c| c.eq_ignore_ascii_case(viewer_country)) + .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] @@ -297,11 +289,13 @@ impl TakedownPredicates<'_> { } 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, } } diff --git a/visibility-filtering/rules/fixtures.rs b/visibility-filtering/rules/fixtures.rs index b0c491ed..614bb9ab 100644 --- a/visibility-filtering/rules/fixtures.rs +++ b/visibility-filtering/rules/fixtures.rs @@ -1,14 +1,44 @@ use crate::models::{ - AuthorFeatures, HydratedTweetCandidate, SafetyLabel, SafetyLabelMap, SafetyLabelType, - TweetFeatures, UserLabelSet, Viewer, ViewerAuthorRelationship, ViewerFeatures, + AuthorFeatures, HydratedTweetCandidate, NsfwFeature, SafetyLabelMap, SafetyLabelType, + TweetFeatures, UserLabelSet, VfAction, Viewer, ViewerAuthorRelationship, ViewerFeatures, }; -use std::collections::{HashMap, HashSet}; +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), @@ -41,14 +71,14 @@ pub(crate) fn candidate() -> CandidateBuilder { author_id: AUTHOR_ID, ..Default::default() }, - labels: HashMap::new(), + labels: HashSet::new(), user_labels: HashSet::new(), } } pub(crate) struct CandidateBuilder { candidate: HydratedTweetCandidate, - labels: HashMap, + labels: HashSet, user_labels: HashSet, } @@ -64,7 +94,7 @@ impl CandidateBuilder { } pub(crate) fn with_label(mut self, label: SafetyLabelType) -> Self { - self.labels.insert(label, SafetyLabel::default()); + self.labels.insert(label); self } @@ -114,3 +144,48 @@ impl CandidateBuilder { 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 index d97ed59e..cc4da723 100644 --- a/visibility-filtering/rules/golden_corpus.rs +++ b/visibility-filtering/rules/golden_corpus.rs @@ -26,7 +26,7 @@ struct Case { #[test] fn golden_corpus_pins_policy_verdicts() { - let rule_engine = RuleEngine::new(); + 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); @@ -54,7 +54,7 @@ fn golden_corpus_pins_policy_verdicts() { #[test] fn every_wired_rule_decides_a_corpus_case() { - let rule_engine = RuleEngine::new(); + 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)) @@ -147,9 +147,12 @@ fn stale_candidate() -> HydratedTweetCandidate { } fn takedown_candidate(reason: TakedownReason) -> HydratedTweetCandidate { - let mut features = TweetFeatures::default(); - features.takedown.reasons = vec![reason]; - candidate().with_tweet_features(features).build() + candidate() + .with_tweet_features(TweetFeatures { + takedown_reasons: vec![reason], + ..Default::default() + }) + .build() } fn exclusive_candidate(viewer_super_follows_author: bool) -> HydratedTweetCandidate { @@ -236,6 +239,54 @@ fn baseline_cases() -> Vec { 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, + }, ] } @@ -524,6 +575,52 @@ fn tweet_shape_cases() -> Vec { 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, + }, ] } diff --git a/visibility-filtering/rules/mod.rs b/visibility-filtering/rules/mod.rs index d28d532b..83da0392 100644 --- a/visibility-filtering/rules/mod.rs +++ b/visibility-filtering/rules/mod.rs @@ -38,11 +38,6 @@ pub(crate) fn test_context<'a>( use std::sync::LazyLock; static NSFW_GATING_COUNTRIES: LazyLock = - LazyLock::new(crate::params::NsfwGatingCountries::new); - RuleContext::new( - SafetyLevel::TimelineHome, - viewer, - candidate, - &NSFW_GATING_COUNTRIES, - ) + LazyLock::new(crate::params::NsfwGatingCountries::starting_at_default); + RuleContext::new(viewer, candidate, &NSFW_GATING_COUNTRIES) } diff --git a/visibility-filtering/rules/registry.rs b/visibility-filtering/rules/registry.rs index 41a332bd..af79ac33 100644 --- a/visibility-filtering/rules/registry.rs +++ b/visibility-filtering/rules/registry.rs @@ -109,8 +109,9 @@ pub struct RuleEngine { } impl RuleEngine { - pub fn new() -> Self { - Self::with_nsfw_gating_countries(Arc::new(NsfwGatingCountries::new())) + #[cfg(test)] + pub(crate) fn for_tests() -> Self { + Self::with_nsfw_gating_countries(Arc::new(NsfwGatingCountries::starting_at_default())) } pub fn with_nsfw_gating_countries(gating_countries: Arc) -> Self { @@ -133,7 +134,7 @@ impl RuleEngine { viewer: &ViewerFeatures, candidate: &HydratedTweetCandidate, ) -> Verdict { - let context = RuleContext::new(level, viewer, candidate, &self.nsfw_gating_countries); + let context = RuleContext::new(viewer, candidate, &self.nsfw_gating_countries); Self::select(level).evaluate(&context) } @@ -150,23 +151,15 @@ impl RuleEngine { } } -impl Default for RuleEngine { - fn default() -> Self { - Self::new() - } -} - #[cfg(test)] mod tests { use super::*; - use crate::models::{ - HydratedTweetCandidate, MediaFeature, TweetFeatures, VfAction, ViewerFeatures, - }; - use crate::rules::fixtures::{author_viewer, candidate, viewer, VIEWER_ID}; + use crate::models::{VfAction, ViewerFeatures}; + use crate::rules::fixtures::{candidate, viewer, VIEWER_ID}; #[test] fn refreshed_config_country_reaches_the_wired_rule() { - let gating_countries = Arc::new(NsfwGatingCountries::new()); + 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) @@ -204,7 +197,7 @@ rust_vf: #[test] fn wired_rule_order_matches_pre_migration_sequence() { - let rule_engine = RuleEngine::new(); + let rule_engine = RuleEngine::for_tests(); assert_eq!( rule_engine.wired_rule_names(SafetyLevel::FilterAll), vec!["FilterAllRule"] @@ -278,369 +271,6 @@ rust_vf: ); } - #[test] - fn filter_all_policy_drops_pristine_candidate() { - let rule_engine = RuleEngine::new(); - let candidate = candidate().build(); - let verdict = rule_engine.evaluate( - SafetyLevel::FilterAll, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(verdict.action, VfAction::Drop(_))); - - let verdict = rule_engine.evaluate( - SafetyLevel::TimelineHome, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(verdict.action, VfAction::Allow)); - - let verdict = rule_engine.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(verdict.action, VfAction::Allow)); - } - - #[test] - fn dmca_media_drops_recommendations_only() { - let rule_engine = RuleEngine::new(); - let candidate = candidate() - .with_tweet_features(TweetFeatures { - media: MediaFeature { - has_dmca_media: true, - ..Default::default() - }, - ..Default::default() - }) - .build(); - - let timeline_home = rule_engine.evaluate( - SafetyLevel::TimelineHome, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(timeline_home.action, VfAction::Allow)); - - let recommendations = rule_engine.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &ViewerFeatures::default(), - &candidate, - ); - assert!(matches!(recommendations.action, VfAction::Drop(_))); - } - - #[test] - fn tweet_nsfw_flag_drops_recommendations_only() { - use crate::models::NsfwFeature; - let rule_engine = RuleEngine::new(); - let candidate = candidate() - .with_tweet_features(TweetFeatures { - nsfw: NsfwFeature { - user: true, - admin: false, - }, - ..Default::default() - }) - .build(); - let viewer = viewer(VIEWER_ID); - - let timeline_home = rule_engine - .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 = rule_engine.evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ); - assert!(matches!(recommendations.action, VfAction::Drop(_))); - assert_eq!(recommendations.decided_by, Some("TweetNsfwUserDropRule")); - } - - #[test] - fn nsfw_author_interstitials_in_network_but_drops_oon() { - use crate::models::AuthorFeatures; - let rule_engine = RuleEngine::new(); - let candidate = candidate() - .with_media() - .with_author_features(AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }) - .build(); - let viewer = viewer(VIEWER_ID); - - let in_network = rule_engine - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Interstitial(_)), - "in-network NSFW author should interstitial, got {in_network:?}" - ); - - let oon = rule_engine - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON NSFW author should drop, got {oon:?}" - ); - } - - #[test] - fn egregious_nsfw_does_not_drop() { - use crate::models::SafetyLabelType; - use xai_x_thrift::user_labels::LabelValue; - let rule_engine = RuleEngine::new(); - - let tweet_candidate = candidate() - .with_label(SafetyLabelType::EGREGIOUS_NSFW) - .build(); - let user_candidate = candidate_with_author_user_label(LabelValue::EGREGIOUS_NSFW, false); - let viewer = viewer(VIEWER_ID); - - for candidate in [&tweet_candidate, &user_candidate] { - let in_network = rule_engine - .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 = rule_engine - .evaluate(SafetyLevel::TimelineHomeRecommendations, &viewer, candidate) - .action; - assert!( - matches!(oon, VfAction::Allow), - "OON EgregiousNsfw should allow after rule removal, got {oon:?}" - ); - } - } - - fn fosnr_candidate( - label: crate::models::SafetyLabelType, - follows: bool, - ) -> HydratedTweetCandidate { - let mut c = candidate().with_label(label).build(); - c.relationship.viewer_follows_author = follows; - c - } - - #[test] - fn fosnr_abuse_insults_drops_oon_but_allows_in_network() { - use crate::models::SafetyLabelType; - let rule_engine = RuleEngine::new(); - let viewer = viewer(VIEWER_ID); - let author = author_viewer(); - - for follows in [true, false] { - let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, follows); - let in_network = rule_engine - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network FosnrAbuseInsults should allow (follows={follows}), got {in_network:?}" - ); - } - - let candidate = fosnr_candidate(SafetyLabelType::FOSNR_ABUSE_INSULTS, false); - let oon = rule_engine - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON FosnrAbuseInsults should drop non-author, got {oon:?}" - ); - - let oon_author = rule_engine - .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() { - let rule_engine = RuleEngine::new(); - let candidate = candidate() - .with_tweet_features(TweetFeatures { - media: MediaFeature { - geo_deny_list: vec!["de".to_string()], - ..Default::default() - }, - ..Default::default() - }) - .build(); - let viewer = ViewerFeatures { - country_code: Some("de".to_string()), - ..viewer(VIEWER_ID) - }; - - let in_network = rule_engine - .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 = rule_engine.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::SafetyLabelType; - let rule_engine = RuleEngine::new(); - let candidate = candidate().with_label(SafetyLabelType::NSFW_TEXT).build(); - let viewer = viewer(VIEWER_ID); - - let in_network = rule_engine - .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:?}" - ); - - let oon = rule_engine - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Drop(_)), - "OON NsfwText should drop, got {oon:?}" - ); - } - - fn candidate_with_author_user_label( - label: xai_x_thrift::user_labels::LabelValue, - follows: bool, - ) -> HydratedTweetCandidate { - let mut c = candidate().with_author_user_label(label).build(); - c.relationship.viewer_follows_author = follows; - c - } - - #[test] - fn nsfw_avatar_user_label_drops_oon_but_allows_in_network() { - use xai_x_thrift::user_labels::LabelValue; - let rule_engine = RuleEngine::new(); - let candidate = candidate_with_author_user_label(LabelValue::NSFW_AVATAR_IMAGE, false); - let viewer = viewer(VIEWER_ID); - - let in_network = rule_engine - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network NsfwAvatarImage should allow, got {in_network:?}" - ); - - let oon = rule_engine.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")); - } - - #[test] - fn recommendations_blacklist_does_not_drop() { - use xai_x_thrift::user_labels::LabelValue; - let rule_engine = RuleEngine::new(); - let candidate = - candidate_with_author_user_label(LabelValue::RECOMMENDATIONS_BLACKLIST, false); - let viewer = viewer(VIEWER_ID); - - let in_network = rule_engine - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network RecommendationsBlacklist should allow, got {in_network:?}" - ); - - let oon = rule_engine - .evaluate( - SafetyLevel::TimelineHomeRecommendations, - &viewer, - &candidate, - ) - .action; - assert!( - matches!(oon, VfAction::Allow), - "OON RecommendationsBlacklist should allow after rule removal, got {oon:?}" - ); - } - - #[test] - fn abusive_high_recall_drops_oon_non_follower_but_allows_in_network() { - use xai_x_thrift::user_labels::LabelValue; - let rule_engine = RuleEngine::new(); - let candidate = candidate_with_author_user_label(LabelValue::ABUSIVE_HIGH_RECALL, false); - let viewer = viewer(VIEWER_ID); - - let in_network = rule_engine - .evaluate(SafetyLevel::TimelineHome, &viewer, &candidate) - .action; - assert!( - matches!(in_network, VfAction::Allow), - "in-network AbusiveHighRecall should allow, got {in_network:?}" - ); - - let oon = rule_engine.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")); - } - #[derive(Debug, PartialEq, Eq)] enum RowClass { Drop, diff --git a/visibility-filtering/rules/tweet_rules.rs b/visibility-filtering/rules/tweet_rules.rs index d868a5eb..15db052e 100644 --- a/visibility-filtering/rules/tweet_rules.rs +++ b/visibility-filtering/rules/tweet_rules.rs @@ -335,37 +335,15 @@ mod tests { use super::*; use crate::models::{ AuthorFeatures, ExclusiveContentFeatures, HydratedTweetCandidate, MediaFeature, - NsfwFeature, TakedownFeature, TweetFeatures, VfAction, Viewer, ViewerAge, ViewerFeatures, + NsfwFeature, TweetFeatures, VfAction, Viewer, ViewerAge, ViewerFeatures, }; use crate::rules::fixtures::{ - author_viewer, candidate, logged_out_viewer, sensitive_opt_in_viewer, viewer, VIEWER_ID, + 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, RuleContext}; + use crate::rules::test_context; use xai_core_entities::entities::{EditControl, EditControlInitial, TakedownReason}; - 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() - ); - } - - 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() - ); - } - fn trigger_label(name: &str) -> SafetyLabelType { match name { "PdnaTweetLabelRule" => SafetyLabelType::PDNA, @@ -503,48 +481,6 @@ mod tests { } } - fn custom_drop(_context: &RuleContext<'_>) -> VfAction { - VfAction::Drop(FilteredReason::UnspecifiedReason) - } - - fn custom_allow(_context: &RuleContext<'_>) -> VfAction { - VfAction::Allow - } - - fn custom_drop_even_self_view(context: &RuleContext<'_>) -> VfAction { - if context.tweet().is_nullcast() { - VfAction::Drop(FilteredReason::TweetIsNullcast) - } else { - VfAction::Allow - } - } - - #[test] - fn custom_row_returns_leaf_action() { - let drop_row = RuleSpec::Custom { - name: "CustomDropLeaf", - evaluate: custom_drop, - }; - let allow_row = RuleSpec::Custom { - name: "CustomAllowLeaf", - evaluate: custom_allow, - }; - let pristine = candidate().build(); - assert_drops( - &drop_row, - &viewer(VIEWER_ID), - &pristine, - &FilteredReason::UnspecifiedReason, - ); - assert_drops( - &drop_row, - &author_viewer(), - &pristine, - &FilteredReason::UnspecifiedReason, - ); - assert_allows(&allow_row, &viewer(VIEWER_ID), &pristine); - } - fn exclusive_candidate( tweet_id: u64, author_id: u64, @@ -558,39 +494,6 @@ mod tests { c } - fn nsfw_flag_media_candidates() -> Vec { - let author_user = candidate() - .with_media() - .with_author_features(AuthorFeatures { - is_nsfw_user: true, - ..Default::default() - }) - .build(); - let tweet_user = { - let mut c = author_user.clone(); - c.author_features = AuthorFeatures::default(); - c.tweet_features.nsfw = NsfwFeature { - user: true, - admin: false, - }; - c - }; - let tweet_admin = { - let mut c = tweet_user.clone(); - c.tweet_features.nsfw = NsfwFeature { - user: false, - admin: true, - }; - c - }; - let both = { - let mut c = tweet_user.clone(); - c.author_features.is_nsfw_admin = true; - c - }; - vec![author_user, tweet_user, tweet_admin, both] - } - #[test] fn nsfw_author_interstitial_axis() { let spec = &NSFW_AUTHOR_INTERSTITIAL[0]; @@ -613,10 +516,10 @@ mod tests { assert_allows(spec, &sensitive_opt_in_viewer(), &firing); assert_allows(spec, &author_viewer(), &firing); } - let mut no_media = nsfw_flag_media_candidates().remove(0); + 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().remove(1); + let [_, mut no_flags, ..] = nsfw_flag_media_candidates(); no_flags.tweet_features.nsfw = NsfwFeature::default(); assert_allows(spec, &viewer(VIEWER_ID), &no_flags); } @@ -841,10 +744,7 @@ mod tests { fn takedown_candidate(reasons: Vec) -> HydratedTweetCandidate { candidate() .with_tweet_features(TweetFeatures { - takedown: TakedownFeature { - reasons, - ..Default::default() - }, + takedown_reasons: reasons, ..Default::default() }) .build() @@ -960,7 +860,6 @@ mod tests { assert_allows(local, &viewer_with_country("de"), &author_local); let non_country = takedown_candidate(vec![ - TakedownReason::Dmca, TakedownReason::HatefulImagery, TakedownReason::Unknown, ]); @@ -968,6 +867,49 @@ mod tests { 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"); @@ -1102,27 +1044,6 @@ mod tests { assert_drops(no_age, &request_fallback, &hp, &reason); } - #[test] - fn custom_row_does_not_add_author_exemption() { - let spec = RuleSpec::Custom { - name: "CustomNullcastLeaf", - evaluate: custom_drop_even_self_view, - }; - let firing = candidate() - .with_tweet_features(TweetFeatures { - is_nullcast: true, - ..Default::default() - }) - .build(); - assert_drops( - &spec, - &author_viewer(), - &firing, - &FilteredReason::TweetIsNullcast, - ); - assert_allows(&spec, &author_viewer(), &candidate().build()); - } - 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, 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 b1c4163a..fcabb52a 100644 --- a/visibility-filtering/safety_label_source/lookup.rs +++ b/visibility-filtering/safety_label_source/lookup.rs @@ -349,9 +349,9 @@ mod tests { #[tokio::test] async fn full_warm_channel_does_not_affect_fallback_result() { - use super::super::warmer::SampledWarmer; + use super::super::warmer::CacheWarmer; - let (warmer, _rx) = SampledWarmer::without_drain_task(1, 100); + let (warmer, _rx) = CacheWarmer::without_drain_task(1); let warmer = Arc::new(warmer); warmer.warm(vec![0]); 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 695b65db..4e53d03c 100644 --- a/visibility-filtering/safety_label_source/metrics.rs +++ b/visibility-filtering/safety_label_source/metrics.rs @@ -206,7 +206,6 @@ pub(crate) fn record_cache_fallback_keys( #[derive(Clone, Copy)] pub(crate) enum WarmKeyResult { EligibleMiss, - SampledOut, Enqueued, DroppedChannelFull, FetchIssued, @@ -217,7 +216,6 @@ impl WarmKeyResult { fn as_str(self) -> &'static str { match self { Self::EligibleMiss => "eligible_miss", - Self::SampledOut => "sampled_out", Self::Enqueued => "enqueued", Self::DroppedChannelFull => "dropped_channel_full", Self::FetchIssued => "fetch_issued", diff --git a/visibility-filtering/safety_label_source/mod.rs b/visibility-filtering/safety_label_source/mod.rs index 008d404f..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; 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 index 18146efa..c5047500 100644 --- a/visibility-filtering/safety_label_source/warmer.rs +++ b/visibility-filtering/safety_label_source/warmer.rs @@ -53,36 +53,27 @@ impl WarmFetcher for StratoWarmFetcher { } } -pub(crate) struct SampledWarmer { +pub(crate) struct CacheWarmer { tx: mpsc::Sender>, - sample_pct: u8, } -impl SampledWarmer { - pub(crate) fn spawn(fetcher: Arc, sample_pct: u8) -> Arc { +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, sample_pct }) + Arc::new(Self { tx }) } #[cfg(test)] - pub(crate) fn without_drain_task( - capacity: usize, - sample_pct: u8, - ) -> (Self, mpsc::Receiver>) { + pub(crate) fn without_drain_task(capacity: usize) -> (Self, mpsc::Receiver>) { let (tx, rx) = mpsc::channel(capacity); - (Self { tx, sample_pct }, rx) + (Self { tx }, rx) } } -impl Warmer for SampledWarmer { - fn warm(&self, mut miss_ids: Vec) { +impl Warmer for CacheWarmer { + fn warm(&self, miss_ids: Vec) { metrics::record_cache_warm_keys(WarmKeyResult::EligibleMiss, miss_ids.len()); - if self.sample_pct < 100 { - let eligible = miss_ids.len(); - miss_ids.retain(|_| fastrand::u8(..100) < self.sample_pct); - metrics::record_cache_warm_keys(WarmKeyResult::SampledOut, eligible - miss_ids.len()); - } if miss_ids.is_empty() { return; } @@ -145,7 +136,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn drain_lingers_then_flushes_in_chunks() { let fetcher = FakeFetcher::new(0); - let warmer = SampledWarmer::spawn(fetcher.clone(), 100); + let warmer = CacheWarmer::spawn(fetcher.clone()); warmer.warm((0..30).collect()); warmer.warm((30..60).collect()); @@ -163,7 +154,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn accumulation_is_capped_per_flush() { let fetcher = FakeFetcher::new(0); - let warmer = SampledWarmer::spawn(fetcher.clone(), 100); + let warmer = CacheWarmer::spawn(fetcher.clone()); for start in (0..150).step_by(30) { warmer.warm((start..start + 30).collect()); @@ -180,7 +171,7 @@ mod tests { #[tokio::test(start_paused = true)] async fn fetch_failure_does_not_stop_the_drain() { let fetcher = FakeFetcher::new(1); - let warmer = SampledWarmer::spawn(fetcher.clone(), 100); + let warmer = CacheWarmer::spawn(fetcher.clone()); warmer.warm(vec![1]); tokio::time::sleep(WARM_LINGER * 2).await; @@ -192,7 +183,7 @@ mod tests { #[tokio::test] async fn full_channel_drops_without_blocking() { - let (warmer, mut rx) = SampledWarmer::without_drain_task(1, 100); + let (warmer, mut rx) = CacheWarmer::without_drain_task(1); warmer.warm(vec![1]); warmer.warm(vec![2]); diff --git a/visibility-filtering/server_deps.rs b/visibility-filtering/server_deps.rs index 4e7501af..12c8f579 100644 --- a/visibility-filtering/server_deps.rs +++ b/visibility-filtering/server_deps.rs @@ -2,14 +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::{SampledWarmer, StratoWarmFetcher, Warmer}; +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; @@ -85,6 +85,10 @@ where } } +#[expect( + clippy::expect_used, + reason = "startup fail-fast: init failure is fatal" +)] pub async fn build_prod_server( datacenter: &str, feature_switches: Arc, @@ -94,14 +98,9 @@ pub async fn build_prod_server( 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 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, @@ -208,9 +207,9 @@ pub async fn build_prod_server( gizmoduck_client, sg_client, safety_label_source.clone(), - fallback_cache_mode, + fallback_cache, ); - let gating_countries = Arc::new(crate::params::NsfwGatingCountries::new()); + 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); @@ -222,7 +221,7 @@ pub async fn build_prod_server( info!( hydrator_count = 5, - ?fallback_cache_mode, + fallback_cache_enabled, home_rule_count, recommendations_rule_count, "VFServer initialized with prod clients" @@ -238,6 +237,7 @@ 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(), @@ -251,6 +251,10 @@ async fn build_reference_compare_harness( "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(); @@ -274,12 +278,15 @@ async fn build_reference_compare_harness( 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> { - let sample_pct = crate::config::cache_warm_sample_pct(); - if sample_pct == 0 { + if !crate::config::cache_warm_enabled() { return None; } @@ -308,11 +315,8 @@ async fn build_cache_warmer( }) .await .expect("Failed to initialize Strato cache-warm client"); - info!(sample_pct, "L2 cache warmer enabled"); - Some(SampledWarmer::spawn( - Arc::new(StratoWarmFetcher::new(grpc)), - sample_pct, - )) + info!("L2 cache warmer enabled"); + Some(CacheWarmer::spawn(Arc::new(StratoWarmFetcher::new(grpc)))) } const TES_STRATO_REQUEST_TIMEOUT_MS: u64 = 100; @@ -445,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; @@ -555,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/connection.rs b/visibility-filtering/twemcache/connection.rs index 975f4e99..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:?}")))?; diff --git a/visibility-filtering/twemcache/host_pool.rs b/visibility-filtering/twemcache/host_pool.rs index 555f42fa..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}; @@ -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; } @@ -221,7 +229,7 @@ impl HostPool { 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; } @@ -238,7 +246,7 @@ impl HostPool { 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 @@ -283,7 +291,7 @@ impl HostPool { 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; @@ -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 { @@ -333,7 +341,7 @@ impl HostPool { 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/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()); } } From 902a06fd616ed815f660e5546d16d492fa1ca825 Mon Sep 17 00:00:00 2001 From: CI agent Date: Fri, 4 Sep 2026 21:03:46 +0000 Subject: [PATCH 17/18] Open-source X Recommendation Algorithm --- .../engagement_counts_hydrator.rs | 17 +- home-mixer/models/candidate.rs | 8 + home-mixer/params/param.rs | 2 +- home-mixer/scorers/author_cold_start.rs | 275 ++++++++++++++---- .../mutual_follow_stats_side_effect.rs | 9 +- .../xai-recsys-proto/proto/recsys.proto | 4 + .../common/xai-proto/proto/recsys.proto | 4 + .../xai_checkpointing/dek.py | 47 ++- .../xai_checkpointing/orbax_encrypted.py | 2 + phoenix/xrex/data/retrieval_dataset.py | 47 +-- phoenix/xrex/train/recsys_bundle_export.py | 10 +- phoenix/xrex/train/trainer_recsys.py | 7 +- phoenix/xrex/utils/log_util.py | 38 +++ visibility-filtering/config.rs | 12 +- visibility-filtering/dark_traffic_setup.rs | 115 +++++--- visibility-filtering/hydration/mod.rs | 4 +- .../hydration/tes_hydrator.rs | 21 ++ visibility-filtering/lib.rs | 12 +- visibility-filtering/models/mod.rs | 2 +- visibility-filtering/models/safety_labels.rs | 1 + visibility-filtering/models/tweet.rs | 4 + visibility-filtering/params.rs | 6 +- visibility-filtering/server.rs | 4 +- visibility-filtering/twemcache/client.rs | 4 - visibility-filtering/twemcache/mod.rs | 4 +- 25 files changed, 480 insertions(+), 179 deletions(-) diff --git a/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs b/home-mixer/candidate_hydrators/engagement_counts_hydrator.rs index b87695f6..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; @@ -86,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 { @@ -218,11 +216,12 @@ 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] diff --git a/home-mixer/models/candidate.rs b/home-mixer/models/candidate.rs index f94d4a4b..a9d747d6 100644 --- a/home-mixer/models/candidate.rs +++ b/home-mixer/models/candidate.rs @@ -104,6 +104,10 @@ pub struct SlateContext { pub recon_count_above: Option, #[serde(default)] pub recon_gap_above: Option, + #[serde(default)] + pub exact_k: Option, + #[serde(default)] + pub exact_gap: Option, } impl From for SlateContext { @@ -124,6 +128,8 @@ impl From for SlateContext { 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, } } } @@ -210,6 +216,8 @@ impl CandidateHelpers for PostCandidate { 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, diff --git a/home-mixer/params/param.rs b/home-mixer/params/param.rs index 7fd5eb54..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-09-03T16:23:24Z +// mirrored from config feature-switch defaults; last sync 2026-09-04T16:22:24Z use xai_feature_switches::param; param!( diff --git a/home-mixer/scorers/author_cold_start.rs b/home-mixer/scorers/author_cold_start.rs index 06c2f7f7..a00202ca 100644 --- a/home-mixer/scorers/author_cold_start.rs +++ b/home-mixer/scorers/author_cold_start.rs @@ -41,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) { @@ -144,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() @@ -174,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; } @@ -189,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), } } @@ -250,43 +288,39 @@ fn pick_thompson( } 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 use_ts = query.params.get(EnableColdStartThompsonSampling); + 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 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_on_home.is_some_and(|imp| imp < threshold) + && c.view_count_on_home + .is_some_and(|imp| imp < params.impression_threshold) }) .map(|(i, _)| i) .collect(); - let best_idx = if use_ts { + let best_idx = if params.use_thompson_sampling { pick_thompson( &eligible, candidates, scores, - query.params.get(ColdStartBetaAlpha0), - query.params.get(ColdStartBetaBeta0), - query.params.get(ColdStartImpressionScale), - query.params.get(ColdStartTsTopK) as usize, + params.beta_alpha0, + params.beta_beta0, + params.impression_scale, + params.ts_top_k, &mut rand::rng(), ) } else { @@ -315,24 +349,24 @@ impl AuthorColdStart { candidates: &[PostCandidate], scores: &[f64], ) -> Vec { - let tracked_ids: String = query.params.get(ColdStartTrackedIds); - record_tracked_ids(candidates, &tracked_ids); - - 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 } @@ -345,7 +379,8 @@ mod tests { 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)] @@ -380,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" @@ -397,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(); @@ -408,7 +453,7 @@ rust_home_mixer: FeatureSwitches::with_options( features, Arc::new(ArmChooser { treatment, control }), - Arc::new(NullBucketImpressor::new()), + impressor, None, false, ) @@ -452,13 +497,18 @@ rust_home_mixer: } 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_on_home: Some(view_count_on_home), served_type: Some(pb::ServedType::ForYouPhoenixRetrievalMoe), - ..Default::default() + ..cold_start_candidate_with_favs(author_id, age, view_count_on_home, fav_count) } } @@ -578,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), @@ -588,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), @@ -651,8 +701,8 @@ rust_home_mixer: 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![ - cold_start_candidate_with_favs(1, minutes(10), 0, 0), - cold_start_candidate_with_favs(2, minutes(20), 3, 0), + 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]); @@ -662,8 +712,8 @@ rust_home_mixer: fn treatment_ts_among_top_k_picks_highest_score() { let author_cold_start = cold_start_with_arms(vec![1, 2], vec![]); let candidates = vec![ - cold_start_candidate_with_favs(1, minutes(10), 0, 0), - cold_start_candidate_with_favs(2, minutes(20), 0, 0), + 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]); @@ -680,6 +730,111 @@ rust_home_mixer: 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"); 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/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto index 6bd6f661..2505571e 100644 --- a/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto +++ b/phoenix/crates/serving/xai-recsys-proto/proto/recsys.proto @@ -116,6 +116,8 @@ message PredictNextActionsRequest { PageDecodeParams pageDecode = 19; bool returnBackboneScores = 20; + + map experiment_overrides = 21; } message PageDecodeParams { @@ -1243,6 +1245,8 @@ message SlateContext { optional uint32 reconCosMilli = 13; optional uint32 reconCountAbove = 14; optional uint32 reconGapAbove = 15; + optional uint32 exactK = 16; + optional uint32 exactGap = 17; } message ActionInfo { diff --git a/phoenix/python/common/xai-proto/proto/recsys.proto b/phoenix/python/common/xai-proto/proto/recsys.proto index 6bd6f661..2505571e 100644 --- a/phoenix/python/common/xai-proto/proto/recsys.proto +++ b/phoenix/python/common/xai-proto/proto/recsys.proto @@ -116,6 +116,8 @@ message PredictNextActionsRequest { PageDecodeParams pageDecode = 19; bool returnBackboneScores = 20; + + map experiment_overrides = 21; } message PageDecodeParams { @@ -1243,6 +1245,8 @@ message SlateContext { optional uint32 reconCosMilli = 13; optional uint32 reconCountAbove = 14; optional uint32 reconGapAbove = 15; + optional uint32 exactK = 16; + optional uint32 exactGap = 17; } message ActionInfo { diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py index 4f49bab4..de12f3fc 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/dek.py @@ -4,6 +4,7 @@ import logging import os import pathlib +import random import tempfile import time @@ -13,6 +14,37 @@ 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: @@ -60,7 +92,7 @@ def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict, ) rank_logger.debug("waiting for wrapped DEK at %s", dek_path) time.sleep(0.1) - raw = bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, dek_path)) + raw = _unwrap_with_backoff(kms_client, dek_path) return raw, entry["wrapped"], entry.get("context") or {}, entry @@ -70,12 +102,10 @@ def publish_tree_dek(path: pathlib.Path, kms_client) -> tuple[bytes, str, dict, def adopt_tree_dek(path: pathlib.Path, kms_client) -> bytes: - import xai_kms - dek_path = path / TREE_DEK_NAME if dek_path.exists(): entry = _read_wrapped_dek(dek_path) or {} - raw = bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, dek_path)) + raw = _unwrap_with_backoff(kms_client, dek_path) key_id = entry.get("key_id") else: raw = _unwrap_header_master(kms_client, path) @@ -85,15 +115,13 @@ def adopt_tree_dek(path: pathlib.Path, kms_client) -> bytes: def _unwrap_header_master(kms_client, path: pathlib.Path) -> bytes: - import xai_kms - 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 bytes(xai_kms.nfs.unwrap_shared_dek(kms_client, tmp)) + return _unwrap_with_backoff(kms_client, tmp) finally: os.unlink(tmp) @@ -110,8 +138,9 @@ def _header_wrapped_master(path: pathlib.Path) -> tuple[str, dict]: context = json.loads(head[start + wrapped_len : start + wrapped_len + context_len]) if not field.startswith(_DERIVED_PREFIX): raise ValueError( - f"envelope at {envelope} carries a legacy per-file wrapped key, " - "not the xai-dek1 derived form" + 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: diff --git a/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py index 4bcc91ef..be367773 100644 --- a/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py +++ b/phoenix/python/training/xai-checkpointing/xai_checkpointing/orbax_encrypted.py @@ -330,6 +330,8 @@ def get_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"): diff --git a/phoenix/xrex/data/retrieval_dataset.py b/phoenix/xrex/data/retrieval_dataset.py index fb422f91..1e381dce 100644 --- a/phoenix/xrex/data/retrieval_dataset.py +++ b/phoenix/xrex/data/retrieval_dataset.py @@ -50,8 +50,17 @@ def _bridge_o2_env() -> None: _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) @@ -232,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"), @@ -252,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/video_4to14day.parquet"), - _idx("post_sid_v5_256x6_snapshots_backup/video_4to14day.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/train/recsys_bundle_export.py b/phoenix/xrex/train/recsys_bundle_export.py index 2ceb1070..c7993d0e 100644 --- a/phoenix/xrex/train/recsys_bundle_export.py +++ b/phoenix/xrex/train/recsys_bundle_export.py @@ -29,6 +29,12 @@ 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 @@ -306,8 +312,8 @@ def forward_fn(batch: Any, merged_embeddings: jax.Array): ) 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.float32) - cont_preds = candidate_continuous_predictions.astype(jnp.float32) + 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 diff --git a/phoenix/xrex/train/trainer_recsys.py b/phoenix/xrex/train/trainer_recsys.py index fe357236..fd517792 100644 --- a/phoenix/xrex/train/trainer_recsys.py +++ b/phoenix/xrex/train/trainer_recsys.py @@ -2858,11 +2858,16 @@ def _maybe_build_stablehlo_bundle(self) -> list | 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(bundle_file.data)) + _write_all_bytes(bundle_path, memoryview(data)) except OSError: rank_logger.exception( "StableHLO bundle write failed; disabling for the rest of this run " 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/visibility-filtering/config.rs b/visibility-filtering/config.rs index 55c5466e..58fb1016 100644 --- a/visibility-filtering/config.rs +++ b/visibility-filtering/config.rs @@ -6,20 +6,20 @@ 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"; -pub const DEFAULT_FS_PATH: &str = "/usr/local/config/features/visibility/main/rust_vf.yml"; +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 gizmoduck_client_id() -> String { +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 twemcache_client_name() -> String { +pub(crate) fn twemcache_client_name() -> String { resolve_twemcache_client_name(std::env::var(ENV_TWEMCACHE_CLIENT_NAME).ok().as_deref()) } @@ -36,15 +36,15 @@ pub fn resolve_twemcache_client_name(configured: Option<&str>) -> String { .to_string() } -pub fn dual_call_harness_enabled() -> bool { +pub(crate) fn dual_call_harness_enabled() -> bool { parse_env_flag(std::env::var(ENV_DUAL_CALL_HARNESS_ENABLED).ok().as_deref()) } -pub fn fallback_cache_enabled() -> bool { +pub(crate) fn fallback_cache_enabled() -> bool { parse_env_flag(std::env::var(ENV_FALLBACK_CACHE_ENABLED).ok().as_deref()) } -pub fn cache_warm_enabled() -> bool { +pub(crate) fn cache_warm_enabled() -> bool { parse_env_flag(std::env::var(ENV_CACHE_WARM_ENABLED).ok().as_deref()) } diff --git a/visibility-filtering/dark_traffic_setup.rs b/visibility-filtering/dark_traffic_setup.rs index f37881bf..e6562531 100644 --- a/visibility-filtering/dark_traffic_setup.rs +++ b/visibility-filtering/dark_traffic_setup.rs @@ -20,10 +20,11 @@ use xai_x_rpc::xds_channel_factory::XdsChannelFactory; const CONFIG_PATH: &str = "/config/dark-traffic/dark_traffic.yaml"; -pub const STAGING_NAMESPACE: &str = "visibility"; +pub const MIRROR_NAMESPACE: &str = "visibility"; pub const STAGING_APP_ENV: &str = "staging"; -pub const STAGING_PORT_ID: &str = "grpc"; -pub const STAGING_WORKLOAD_PREFIX: &str = "xai-vf-service"; +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; @@ -34,28 +35,42 @@ 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_staging_listener(listener_name: &str) -> Option { +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 suffix = format!(".{STAGING_APP_ENV}.{STAGING_NAMESPACE}:{STAGING_PORT_ID}"); - let workload = dest.strip_suffix(suffix.as_str())?; - if !workload.starts_with(STAGING_WORKLOAD_PREFIX) || workload.contains('.') { + 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: workload.to_string(), + name, xds_dest: dest.to_string(), }) } -struct XdsStagingDiscovery { +struct XdsMirrorDiscovery { server_uri: String, } #[async_trait::async_trait] -impl EndpointDiscovery for XdsStagingDiscovery { +impl EndpointDiscovery for XdsMirrorDiscovery { async fn discover(&self) -> anyhow::Result> { tokio::time::timeout(LDS_RESPONSE_TIMEOUT, self.fetch()) .await @@ -63,7 +78,7 @@ impl EndpointDiscovery for XdsStagingDiscovery { } } -impl XdsStagingDiscovery { +impl XdsMirrorDiscovery { async fn fetch(&self) -> anyhow::Result> { let channel = tonic::transport::Endpoint::from_shared(self.server_uri.clone()) .context("invalid kube-discovery URI")? @@ -78,8 +93,8 @@ impl XdsStagingDiscovery { type_url: LISTENER_TYPE_URL.to_string(), resource_names: vec!["*".to_string()], node: Some(Node { - id: format!("{STAGING_WORKLOAD_PREFIX}-dark-traffic"), - cluster: STAGING_NAMESPACE.to_string(), + id: format!("{MIRROR_WORKLOAD_PREFIX}-dark-traffic"), + cluster: MIRROR_NAMESPACE.to_string(), ..Default::default() }), ..Default::default() @@ -116,13 +131,13 @@ impl XdsStagingDiscovery { ); let endpoints: Vec = names .iter() - .filter_map(|name| parse_staging_listener(name)) + .filter_map(|name| parse_mirror_listener(name)) .collect(); if endpoints.is_empty() { tracing::warn!( listeners = names.len(), - "dark_traffic: no staging listeners matched" + "dark_traffic: no staging or devel listeners matched" ); } else { info!( @@ -135,14 +150,20 @@ impl XdsStagingDiscovery { } } -struct TimeoutChannelFactory { - inner: XdsChannelFactory, +struct MirrorChannelFactory { + staging: XdsChannelFactory, + devel: XdsChannelFactory, } #[async_trait::async_trait] -impl ChannelFactory for TimeoutChannelFactory { +impl ChannelFactory for MirrorChannelFactory { async fn create_channel(&self, ep: &EndpointInfo) -> anyhow::Result { - tokio::time::timeout(CHANNEL_CREATE_TIMEOUT, self.inner.create_channel(ep)) + 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))? } @@ -175,23 +196,21 @@ pub fn resolve_layer() -> DarkLayer { } let dc = std::env::var("DATACENTER").unwrap_or_else(|_| "atla".to_string()); - let discovery = XdsStagingDiscovery { + let discovery = XdsMirrorDiscovery { server_uri: format!("http://frontend.kube-discovery.prod.svc.{dc}.kube.int-x.ai:8082"), }; - let domain = staging_tls_domain(&dc); - info!(domain, "dark_traffic: enabled"); + 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 factory = XdsChannelFactory::new( - TlsMode::mtls_from_env() - .expect("S2S TLS config required") - .with_domain_override(&domain), - ); + 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(TimeoutChannelFactory { inner: factory }), - Arc::new(discovery), - ); + let channels = DynamicChannelManager::new(Arc::new(factory), Arc::new(discovery)); let config = ReloadableDarkTrafficConfigBuilder::new(CONFIG_PATH) .forwarders({ @@ -262,9 +281,21 @@ mod tests { } #[test] - fn parse_accepts_vf_staging_listeners() { + 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", @@ -274,9 +305,18 @@ mod tests { "xai-vf-service", ), ] { - let ep = parse_staging_listener(listener).expect(listener); + let ep = parse_mirror_listener(listener).expect(listener); assert_eq!(ep.name, workload); - assert_eq!(ep.xds_dest, format!("{workload}.staging.visibility:grpc")); + assert_eq!( + ep.xds_dest, + listener + .rsplit("/") + .next() + .unwrap() + .split("?") + .next() + .unwrap() + ); } } @@ -284,6 +324,11 @@ mod tests { 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", @@ -291,7 +336,7 @@ mod tests { "xai-vf-service.staging.visibility", "", ] { - assert!(parse_staging_listener(name).is_none(), "{name}"); + assert!(parse_mirror_listener(name).is_none(), "{name}"); } } } diff --git a/visibility-filtering/hydration/mod.rs b/visibility-filtering/hydration/mod.rs index 7f905586..0ea33a16 100644 --- a/visibility-filtering/hydration/mod.rs +++ b/visibility-filtering/hydration/mod.rs @@ -264,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() @@ -304,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/tes_hydrator.rs b/visibility-filtering/hydration/tes_hydrator.rs index fb5f21c5..7b18980a 100644 --- a/visibility-filtering/hydration/tes_hydrator.rs +++ b/visibility-filtering/hydration/tes_hydrator.rs @@ -432,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( diff --git a/visibility-filtering/lib.rs b/visibility-filtering/lib.rs index 4e3e44b0..bef575e7 100644 --- a/visibility-filtering/lib.rs +++ b/visibility-filtering/lib.rs @@ -22,18 +22,18 @@ ) )] -pub mod clients; +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(crate) mod hydration; +pub(crate) mod models; pub mod params; pub(crate) mod reference_compare; -pub mod rules; -pub mod safety_label_source; +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/models/mod.rs b/visibility-filtering/models/mod.rs index 237aaacd..b00cd656 100644 --- a/visibility-filtering/models/mod.rs +++ b/visibility-filtering/models/mod.rs @@ -10,7 +10,7 @@ pub use exclusive_content::ExclusiveContentFeatures; pub use relationship::ViewerAuthorRelationship; pub use safety_labels::{SafetyLabelMap, SafetyLabelType}; pub use tweet::{CoreFeature, MediaFeature, NsfwFeature, TweetFeatures}; -pub use viewer::{Viewer, ViewerAge, ViewerFeatures, ADULT_AGE_YEARS}; +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 2ef75f06..5389ea5d 100644 --- a/visibility-filtering/models/safety_labels.rs +++ b/visibility-filtering/models/safety_labels.rs @@ -7,6 +7,7 @@ use xai_visibility_filtering_proto as vf_pb; pub struct SafetyLabelMap(HashSet); impl SafetyLabelMap { + #[cfg(test)] pub fn new(label_types: HashSet) -> Self { Self(label_types) } diff --git a/visibility-filtering/models/tweet.rs b/visibility-filtering/models/tweet.rs index f6f6a922..37098848 100644 --- a/visibility-filtering/models/tweet.rs +++ b/visibility-filtering/models/tweet.rs @@ -1,5 +1,9 @@ #[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, } diff --git a/visibility-filtering/params.rs b/visibility-filtering/params.rs index 453929de..15598882 100644 --- a/visibility-filtering/params.rs +++ b/visibility-filtering/params.rs @@ -7,8 +7,8 @@ use xai_feature_switches::{FeatureSwitches, RecipientBuilder, Value}; pub const NSFW_GATING_COUNTRIES_KEY: &str = "rust_vf_nsfw_gating_countries"; -pub const SCALA_NSFW_GATING_FILE: &str = "country_specific_nsfw_content_gating.yml"; -pub const SCALA_NSFW_GATING_COUNTRIES_KEY: &str = "country_specific_nsfw_content_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"; @@ -21,7 +21,7 @@ pub fn default_nsfw_gating_countries() -> Vec { .to_vec() } -pub struct NsfwGatingCountries { +pub(crate) struct NsfwGatingCountries { countries: ArcSwap>, } diff --git a/visibility-filtering/server.rs b/visibility-filtering/server.rs index 1ef69856..1b58b4b1 100644 --- a/visibility-filtering/server.rs +++ b/visibility-filtering/server.rs @@ -28,14 +28,14 @@ impl xai_x_service_builder::XService for VFServer { } impl VFServer { - pub async fn new( + 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/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/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}; From 843f3a96c860d3d1238683e53798baac6b628bad Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 8 Sep 2026 02:35:03 +0000 Subject: [PATCH 18/18] Stop SubscriptionHydrator from skipping exclusive ids on quotes Phoenix TES exclusive lookup keyed only the quote wrapper tweet_id. Super Follow lives on the quoted original, so IneligibleSubscriptionFilter kept the card as public. Fetch the quoted id and use it when the wrapper is not exclusive. Retweet originals stay #155. Co-authored-by: Jon Bailey --- .../subscription_hydrator.rs | 172 +++++++++++++++--- 1 file changed, 151 insertions(+), 21 deletions(-) diff --git a/home-mixer/candidate_hydrators/subscription_hydrator.rs b/home-mixer/candidate_hydrators/subscription_hydrator.rs index b5309eb3..e440be40 100644 --- a/home-mixer/candidate_hydrators/subscription_hydrator.rs +++ b/home-mixer/candidate_hydrators/subscription_hydrator.rs @@ -1,6 +1,7 @@ use crate::clients::tweet_entity_service_client::TESClient; use crate::models::candidate::PostCandidate; use crate::models::query::ScoredPostsQuery; +use std::collections::HashMap; use std::sync::Arc; use tonic::async_trait; use xai_candidate_pipeline::component_library::utils::{default_quick_cache, QuickCache}; @@ -52,31 +53,160 @@ impl CachedHydrator for SubscriptionHydrator { ) -> Vec> { let client = &self.tes_client; - let tweet_ids: Vec = candidates.iter().map(|c| c.tweet_id).collect(); - - let post_features = client.get_subscription_author_ids(tweet_ids.clone()).await; - - let mut hydrated_candidates = Vec::with_capacity(candidates.len()); - for tweet_id in tweet_ids { - let post_features = post_features.get(&tweet_id); - let hydrated = match post_features { - Some(Ok(value)) => Ok(PostCandidate { - subscription_author_id: *value, - ..Default::default() - }), - None => Err(format!( - "Missing subscription author id for tweet_id={}", - tweet_id - )), - Some(Err(err)) => Err(err.to_string()), - }; - hydrated_candidates.push(hydrated); - } + let tweet_ids = subscription_fetch_ids(candidates); + + let post_features = client.get_subscription_author_ids(tweet_ids).await; - hydrated_candidates + candidates + .iter() + .map(|candidate| { + resolve_subscription_author_id(&post_features, candidate).map(|value| { + PostCandidate { + subscription_author_id: value, + ..Default::default() + } + }) + }) + .collect() } fn update(&self, candidate: &mut PostCandidate, hydrated: PostCandidate) { candidate.subscription_author_id = hydrated.subscription_author_id; } } + +fn subscription_fetch_ids(candidates: &[PostCandidate]) -> Vec { + let mut ids = Vec::with_capacity(candidates.len() * 2); + for candidate in candidates { + ids.push(candidate.tweet_id); + if let Some(quoted_id) = candidate.quoted_tweet_id { + if quoted_id != candidate.tweet_id { + ids.push(quoted_id); + } + } + } + ids +} + +fn resolve_subscription_author_id( + tes: &HashMap, E>>, + candidate: &PostCandidate, +) -> Result, String> { + let own = match tes.get(&candidate.tweet_id) { + Some(Ok(value)) => *value, + None => { + return Err(format!( + "Missing subscription author id for tweet_id={}", + candidate.tweet_id + )); + } + Some(Err(err)) => return Err(err.to_string()), + }; + + if own.is_some() { + return Ok(own); + } + + if let Some(quoted_id) = candidate.quoted_tweet_id { + if let Some(Ok(Some(id))) = tes.get("ed_id) { + return Ok(Some(*id)); + } + } + + Ok(None) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn native_exclusive_still_uses_candidate_id() { + let candidates = vec![PostCandidate { + tweet_id: 20, + ..Default::default() + }]; + assert_eq!(subscription_fetch_ids(&candidates), vec![20]); + } + + #[test] + fn quote_of_exclusive_fetches_quoted_id() { + let candidates = vec![PostCandidate { + tweet_id: 10, + quoted_tweet_id: Some(20), + ..Default::default() + }]; + assert_eq!(subscription_fetch_ids(&candidates), vec![10, 20]); + } + + #[test] + fn retweet_without_quote_does_not_use_original_id() { + let candidates = vec![PostCandidate { + tweet_id: 10, + retweeted_tweet_id: Some(20), + ..Default::default() + }]; + assert_eq!(subscription_fetch_ids(&candidates), vec![10]); + } + + #[test] + fn tes_keyed_only_by_wrapper_does_not_mark_quote_exclusive() { + let candidate = PostCandidate { + tweet_id: 10, + quoted_tweet_id: Some(20), + ..Default::default() + }; + let mut tes = HashMap::new(); + tes.insert(10, Ok(None)); + assert_eq!( + resolve_subscription_author_id(&tes, &candidate).unwrap(), + None + ); + } + + #[test] + fn quote_of_exclusive_uses_quoted_author() { + let candidate = PostCandidate { + tweet_id: 10, + quoted_tweet_id: Some(20), + ..Default::default() + }; + let mut tes = HashMap::new(); + tes.insert(10, Ok(None)); + tes.insert(20, Ok(Some(99))); + assert_eq!( + resolve_subscription_author_id(&tes, &candidate).unwrap(), + Some(99) + ); + } + + #[test] + fn exclusive_quote_of_public_keeps_wrapper_author() { + let candidate = PostCandidate { + tweet_id: 10, + quoted_tweet_id: Some(20), + ..Default::default() + }; + let mut tes = HashMap::new(); + tes.insert(10, Ok(Some(7))); + tes.insert(20, Ok(None)); + assert_eq!( + resolve_subscription_author_id(&tes, &candidate).unwrap(), + Some(7) + ); + } + + #[test] + fn native_not_exclusive_stays_none() { + let candidate = PostCandidate { + tweet_id: 20, + ..Default::default() + }; + let mut tes = HashMap::new(); + tes.insert(20, Ok(None)); + assert_eq!( + resolve_subscription_author_id(&tes, &candidate).unwrap(), + None + ); + } +}