From eef3547c375f0117a559d5712227620407f055f4 Mon Sep 17 00:00:00 2001 From: Ilia Alshanetsky Date: Thu, 9 Apr 2026 19:00:30 -0400 Subject: [PATCH] perf: optimize dedup, parallelize handle searches and enrichment The dedup hot path recomputed normalize_text() 4 times per comparison and recomputed item_text() on every inner-loop iteration. Pre-computing n-gram sets and token sets into a _PreparedText cache cuts dedup time by 6x (2.16s to 0.39s on 300 unique items). Bird handle searches spawned one Node process per handle sequentially. Now uses ThreadPoolExecutor so N handles run concurrently. Same pattern applied to YouTube comment enrichment (was serial, Reddit was already parallel) and the retry-thin-sources phase in the pipeline. Clustering now pre-computes candidate text and uses prepared_similarity for the O(n^2) grouping and MMR representative selection loops. Minor: _is_wsl() cached with lru_cache, Bundle.add_items() uses extend() instead of list concatenation. End-to-end: 5.2s -> 3.7s (29% faster) on a typical 4-source query. --- scripts/lib/bird_x.py | 21 ++++++---- scripts/lib/cluster.py | 30 ++++++++++---- scripts/lib/cookie_extract.py | 7 +++- scripts/lib/dedupe.py | 32 ++++++++++++++- scripts/lib/pipeline.py | 77 +++++++++++++++++++---------------- scripts/lib/schema.py | 3 +- scripts/lib/youtube_yt.py | 17 ++++++-- 7 files changed, 128 insertions(+), 59 deletions(-) diff --git a/scripts/lib/bird_x.py b/scripts/lib/bird_x.py index 66f3101..9ab906d 100644 --- a/scripts/lib/bird_x.py +++ b/scripts/lib/bird_x.py @@ -312,10 +312,9 @@ def search_handles( Returns: List of raw item dicts (same format as parse_bird_response output). """ - all_items = [] core_topic = _extract_core_subject(topic) if topic else None - for handle in handles: + def _search_one_handle(handle: str) -> List[Dict[str, Any]]: handle = handle.lstrip("@") if core_topic: query = f"from:{handle} {core_topic} since:{from_date}" @@ -350,24 +349,32 @@ def search_handles( proc.kill() proc.wait(timeout=5) _log(f"Handle search timed out for @{handle}") - continue + return [] if proc.returncode != 0: _log(f"Handle search failed for @{handle}: {(stderr or '').strip()}") - continue + return [] output = (stdout or "").strip() if not output: - continue + return [] response = json.loads(output) - items = parse_bird_response(response, query=core_topic) - all_items.extend(items) + return parse_bird_response(response, query=core_topic) except json.JSONDecodeError: _log(f"Invalid JSON from handle search for @{handle}") except (OSError, subprocess.SubprocessError) as e: _log(f"Handle search error for @{handle}: {e}") + return [] + + from concurrent.futures import ThreadPoolExecutor, as_completed + + all_items: List[Dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=min(5, len(handles))) as executor: + futures = {executor.submit(_search_one_handle, h): h for h in handles} + for future in as_completed(futures): + all_items.extend(future.result()) return all_items diff --git a/scripts/lib/cluster.py b/scripts/lib/cluster.py index 60b1043..bd9e501 100644 --- a/scripts/lib/cluster.py +++ b/scripts/lib/cluster.py @@ -57,28 +57,34 @@ def _entity_overlap(entities_a: set[str], entities_b: set[str]) -> float: def _mmr_representatives( candidates: list[schema.Candidate], + text_cache: dict[str, dedupe._PreparedText], limit: int = 3, diversity_lambda: float = 0.75, ) -> list[str]: selected: list[schema.Candidate] = [] + remaining_set = {c.candidate_id for c in candidates} remaining = list(candidates) while remaining and len(selected) < limit: if not selected: best = max(remaining, key=lambda candidate: candidate.final_score) selected.append(best) - remaining.remove(best) + remaining_set.discard(best.candidate_id) + remaining = [c for c in remaining if c.candidate_id in remaining_set] continue + selected_preps = [text_cache[c.candidate_id] for c in selected] + def score(candidate: schema.Candidate) -> float: + prep = text_cache[candidate.candidate_id] diversity_penalty = max( - dedupe.hybrid_similarity(_candidate_text(candidate), _candidate_text(existing)) - for existing in selected + dedupe.prepared_similarity(prep, sp) for sp in selected_preps ) return (diversity_lambda * candidate.final_score) - ((1 - diversity_lambda) * diversity_penalty * 100) best = max(remaining, key=score) selected.append(best) - remaining.remove(best) + remaining_set.discard(best.candidate_id) + remaining = [c for c in remaining if c.candidate_id in remaining_set] return [candidate.candidate_id for candidate in selected] @@ -105,15 +111,21 @@ def cluster_candidates( ) return clusters + text_cache: dict[str, dedupe._PreparedText] = { + c.candidate_id: dedupe._PreparedText(_candidate_text(c)) + for c in candidates + } + groups: list[list[schema.Candidate]] = [] # Lower threshold for breaking_news: related articles share fewer exact # words but cover the same event. threshold = 0.42 if plan.intent == "breaking_news" else 0.48 for candidate in candidates: assigned = False + cand_prep = text_cache[candidate.candidate_id] for group in groups: leader = group[0] - similarity = dedupe.hybrid_similarity(_candidate_text(candidate), _candidate_text(leader)) + similarity = dedupe.prepared_similarity(cand_prep, text_cache[leader.candidate_id]) if similarity >= threshold: group.append(candidate) assigned = True @@ -125,7 +137,7 @@ def cluster_candidates( for index, group in enumerate(groups, start=1): group.sort(key=lambda candidate: candidate.final_score, reverse=True) cluster_id = f"cluster-{index}" - representatives = _mmr_representatives(group) + representatives = _mmr_representatives(group, text_cache) for candidate in group: candidate.cluster_id = cluster_id clusters.append( @@ -225,7 +237,11 @@ def _merge_entity_clusters( # Pick representatives from combined pool combined_candidates = [candidate_map[cid] for cid in combined_cids if cid in candidate_map] combined_candidates.sort(key=lambda c: c.final_score, reverse=True) - reps = _mmr_representatives(combined_candidates) + merge_text_cache = { + c.candidate_id: dedupe._PreparedText(_candidate_text(c)) + for c in combined_candidates + } + reps = _mmr_representatives(combined_candidates, merge_text_cache) cluster_id = cl.cluster_id for cid in combined_cids: diff --git a/scripts/lib/cookie_extract.py b/scripts/lib/cookie_extract.py index 1f2efbf..589c0fe 100644 --- a/scripts/lib/cookie_extract.py +++ b/scripts/lib/cookie_extract.py @@ -7,6 +7,7 @@ Only uses Python stdlib — no external dependencies. """ import configparser +import functools import logging import platform import shutil @@ -18,8 +19,12 @@ from typing import Dict, List, Optional logger = logging.getLogger(__name__) +@functools.lru_cache(maxsize=1) def _is_wsl() -> bool: - """Detect if running under Windows Subsystem for Linux.""" + """Detect if running under Windows Subsystem for Linux. + + Cached after the first call since /proc/version doesn't change at runtime. + """ try: return "microsoft" in Path("/proc/version").read_text().lower() except OSError: diff --git a/scripts/lib/dedupe.py b/scripts/lib/dedupe.py index de95bf9..670340d 100644 --- a/scripts/lib/dedupe.py +++ b/scripts/lib/dedupe.py @@ -76,6 +76,31 @@ def hybrid_similarity(text_a: str, text_b: str) -> float: ) +def _tokenize(normalized: str) -> frozenset[str]: + return frozenset( + tok for tok in normalized.split() + if len(tok) > 1 and tok not in STOPWORDS + ) + + +class _PreparedText: + """Pre-computed text representations for fast repeated similarity checks.""" + + __slots__ = ("ngrams", "tokens") + + def __init__(self, raw: str) -> None: + norm = normalize_text(raw) + self.ngrams = get_ngrams(norm) if norm else set() + self.tokens = _tokenize(norm) + + +def prepared_similarity(a: _PreparedText, b: _PreparedText) -> float: + return max( + jaccard_similarity(a.ngrams, b.ngrams), + jaccard_similarity(a.tokens, b.tokens), + ) + + def item_text(item: schema.SourceItem) -> str: parts = [item.title, item.body, item.author or "", item.container or ""] return " ".join(part for part in parts if part).strip() @@ -84,16 +109,19 @@ def item_text(item: schema.SourceItem) -> str: def dedupe_items(items: list[schema.SourceItem], threshold: float = 0.7) -> list[schema.SourceItem]: """Remove near-duplicates while keeping earlier, better-scored items.""" kept: list[schema.SourceItem] = [] + kept_prepared: list[_PreparedText] = [] for item in items: text = item_text(item) if not text: kept.append(item) continue + prep = _PreparedText(text) is_duplicate = False - for existing in kept: - if hybrid_similarity(text, item_text(existing)) >= threshold: + for existing_prep in kept_prepared: + if prepared_similarity(prep, existing_prep) >= threshold: is_duplicate = True break if not is_duplicate: kept.append(item) + kept_prepared.append(prep) return kept diff --git a/scripts/lib/pipeline.py b/scripts/lib/pipeline.py index 957e404..960f215 100644 --- a/scripts/lib/pipeline.py +++ b/scripts/lib/pipeline.py @@ -719,44 +719,49 @@ def _retry_thin_sources( weight=0.3, ) - for source in thin_sources: - if source in rate_limited_sources: - continue - try: - raw_items, _artifact = _retrieve_stream( - topic=topic, - subquery=retry_subquery, - source=source, - config=config, - depth=depth, - date_range=date_range, - runtime=runtime, - mock=mock, - rate_limited_sources=rate_limited_sources, - rate_limit_lock=rate_limit_lock, - web_backend=web_backend, - raw_topic=topic, - ) - normalized = _normalize_score_dedupe( - source, - raw_items, - from_date, - to_date, - freshness_mode=plan.freshness_mode, - ranking_query=retry_subquery.ranking_query, - ) - normalized = normalized[:settings["per_stream_limit"]] + def _retry_one_source(source: str) -> tuple[str, list[schema.SourceItem]]: + raw_items, _artifact = _retrieve_stream( + topic=topic, + subquery=retry_subquery, + source=source, + config=config, + depth=depth, + date_range=date_range, + runtime=runtime, + mock=mock, + rate_limited_sources=rate_limited_sources, + rate_limit_lock=rate_limit_lock, + web_backend=web_backend, + raw_topic=topic, + ) + normalized = _normalize_score_dedupe( + source, + raw_items, + from_date, + to_date, + freshness_mode=plan.freshness_mode, + ranking_query=retry_subquery.ranking_query, + ) + return source, normalized[:settings["per_stream_limit"]] - existing_urls = {item.url for item in bundle.items_by_source.get(source, []) if item.url} - new_items = [item for item in normalized if item.url not in existing_urls] + retryable = [s for s in thin_sources if s not in rate_limited_sources] - if new_items: - bundle.items_by_source.setdefault(source, []).extend(new_items) - primary_label = plan.subqueries[0].label if plan.subqueries else "primary" - existing = bundle.items_by_source_and_query.get((primary_label, source), []) - bundle.items_by_source_and_query[(primary_label, source)] = existing + new_items - except Exception as exc: - print(f"[Pipeline] Retry failed for {source}: {type(exc).__name__}: {exc}", file=sys.stderr) + from concurrent.futures import ThreadPoolExecutor, as_completed + with ThreadPoolExecutor(max_workers=min(4, len(retryable) or 1)) as executor: + futures = {executor.submit(_retry_one_source, s): s for s in retryable} + for future in as_completed(futures): + source = futures[future] + try: + source, normalized = future.result() + existing_urls = {item.url for item in bundle.items_by_source.get(source, []) if item.url} + new_items = [item for item in normalized if item.url not in existing_urls] + + if new_items: + bundle.items_by_source.setdefault(source, []).extend(new_items) + primary_label = plan.subqueries[0].label if plan.subqueries else "primary" + bundle.items_by_source_and_query.setdefault((primary_label, source), []).extend(new_items) + except Exception as exc: + print(f"[Pipeline] Retry failed for {source}: {type(exc).__name__}: {exc}", file=sys.stderr) def _retrieve_stream( diff --git a/scripts/lib/schema.py b/scripts/lib/schema.py index 20e5be9..d9a600b 100644 --- a/scripts/lib/schema.py +++ b/scripts/lib/schema.py @@ -168,8 +168,7 @@ class RetrievalBundle: def add_items(self, label: str, source: str, items: list[SourceItem]) -> None: """Atomically append items to both items_by_source_and_query and items_by_source.""" - existing = self.items_by_source_and_query.get((label, source), []) - self.items_by_source_and_query[(label, source)] = existing + items + self.items_by_source_and_query.setdefault((label, source), []).extend(items) self.items_by_source.setdefault(source, []).extend(items) diff --git a/scripts/lib/youtube_yt.py b/scripts/lib/youtube_yt.py index d77c3a3..ac26a6e 100644 --- a/scripts/lib/youtube_yt.py +++ b/scripts/lib/youtube_yt.py @@ -699,18 +699,27 @@ def enrich_with_comments( top_items = ranked[:max_videos] _log(f"Enriching comments for {len(top_items)} YouTube videos") - enriched_count = 0 - for item in top_items: + from concurrent.futures import ThreadPoolExecutor, as_completed + + def _enrich_one(item: dict) -> bool: video_id = item.get("video_id", "") if not video_id: - continue + return False try: comments = _fetch_video_comments(video_id, token, max_comments) if comments: item["top_comments"] = comments - enriched_count += 1 + return True except Exception as exc: _log(f"Comment enrichment failed for {video_id}: {exc}") + return False + + enriched_count = 0 + with ThreadPoolExecutor(max_workers=min(4, len(top_items))) as executor: + futures = {executor.submit(_enrich_one, item): item for item in top_items} + for future in as_completed(futures): + if future.result(): + enriched_count += 1 _log(f"Enriched {enriched_count}/{len(top_items)} videos with comments") return items