| 1 | #!/usr/bin/env python3 |
| 2 | """Run a local no-publish GitHub crawl sharding experiment for issue #435.""" |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import argparse |
| 7 | import json |
| 8 | import math |
| 9 | import os |
| 10 | import threading |
| 11 | import time |
| 12 | from collections import Counter |
| 13 | from concurrent.futures import ThreadPoolExecutor, as_completed |
| 14 | from dataclasses import dataclass, field |
| 15 | from datetime import UTC, datetime |
| 16 | from pathlib import Path |
| 17 | from queue import Empty, Queue |
| 18 | from typing import Any |
| 19 | |
| 20 | from scripts.crawl import ( |
| 21 | RAW_ROOT, |
| 22 | GitHubClient, |
| 23 | build_signals, |
| 24 | build_star_snapshot, |
| 25 | collect_repositories, |
| 26 | github_artifact_checksum, |
| 27 | github_crawl_config_checksum, |
| 28 | github_schema_checksum, |
| 29 | iso_timestamp, |
| 30 | load_previous_star_snapshot, |
| 31 | load_topic_queries, |
| 32 | significance_skip_reason, |
| 33 | to_repo_record, |
| 34 | utc_now, |
| 35 | validate_payload, |
| 36 | week_slug, |
| 37 | write_payload, |
| 38 | ) |
| 39 | from scripts.topic_paths import cache_dir, raw_dir, snapshots_dir |
| 40 | |
| 41 | EXPERIMENT_ROOT = Path("data/experiments/shard-435") |
| 42 | DEFAULT_WALL_CLOCK_BUDGET = 120 |
| 43 | DEFAULT_API_BUDGET_MULTIPLIER = 1.1 |
| 44 | |
| 45 | |
| 46 | class ExperimentAbort(RuntimeError): |
| 47 | """Abort the shard experiment immediately.""" |
| 48 | |
| 49 | |
| 50 | class ShardBudgetExceeded(RuntimeError): |
| 51 | """Stop a shard gracefully once its wall-clock budget is exhausted.""" |
| 52 | |
| 53 | |
| 54 | @dataclass(slots=True) |
| 55 | class GuardrailEvent: |
| 56 | kind: str |
| 57 | shard: str |
| 58 | message: str |
| 59 | at: str |
| 60 | details: dict[str, Any] = field(default_factory=dict) |
| 61 | |
| 62 | def to_dict(self) -> dict[str, Any]: |
| 63 | payload = {"kind": self.kind, "shard": self.shard, "message": self.message, "at": self.at} |
| 64 | if self.details: |
| 65 | payload["details"] = self.details |
| 66 | return payload |
| 67 | |
| 68 | |
| 69 | @dataclass(slots=True) |
| 70 | class SearchPlan: |
| 71 | shard_name: str |
| 72 | repo_group: str |
| 73 | primary_queries: list[str] |
| 74 | secondary_queries: list[str] = field(default_factory=list) |
| 75 | min_repos_per_week: int = 0 |
| 76 | |
| 77 | |
| 78 | @dataclass(slots=True) |
| 79 | class ValidationItem: |
| 80 | repo_group: str |
| 81 | sequence: int |
| 82 | repo: dict[str, Any] |
| 83 | |
| 84 | |
| 85 | @dataclass(slots=True) |
| 86 | class ValidatedRecord: |
| 87 | repo_group: str |
| 88 | sequence: int |
| 89 | record: dict[str, Any] |
| 90 | |
| 91 | |
| 92 | @dataclass(slots=True) |
| 93 | class CrawlContext: |
| 94 | args: argparse.Namespace |
| 95 | topic_id: str | None |
| 96 | topic_raw: Path |
| 97 | topic_snapshots: Path |
| 98 | topic_cache: Path |
| 99 | crawled_at: datetime |
| 100 | run_started_at: datetime |
| 101 | since: datetime |
| 102 | window_end: datetime |
| 103 | week: str |
| 104 | max_results: int |
| 105 | config_checksum: str |
| 106 | current_code_sha: str |
| 107 | source_refresh_policy: str |
| 108 | baseline_output_path: Path |
| 109 | baseline_snapshot_path: Path |
| 110 | shard_output_path: Path |
| 111 | shard_snapshot_path: Path |
| 112 | |
| 113 | |
| 114 | @dataclass(slots=True) |
| 115 | class RunResult: |
| 116 | name: str |
| 117 | payload: dict[str, Any] |
| 118 | snapshot_payload: dict[str, Any] |
| 119 | api_calls: int |
| 120 | cache_hits: int |
| 121 | stale_cache_hits: int |
| 122 | rate_limit_events: int |
| 123 | secondary_rate_limit_events: int |
| 124 | partial_failures: list[str] |
| 125 | wall_clock_s: float |
| 126 | shards_used: int |
| 127 | completed: bool |
| 128 | guardrail_events: list[dict[str, Any]] |
| 129 | |
| 130 | |
| 131 | @dataclass(slots=True) |
| 132 | class ShardResult: |
| 133 | shard_id: str |
| 134 | repos_found: list[dict[str, Any]] |
| 135 | api_calls: int |
| 136 | errors: list[str] |
| 137 | wall_clock_s: float |
| 138 | rate_limit_detected: bool = False |
| 139 | rate_limit_events: int = 0 |
| 140 | |
| 141 | |
| 142 | @dataclass(slots=True) |
| 143 | class ExperimentReport: |
| 144 | speedup_pct: float | None |
| 145 | api_growth_pct: float | None |
| 146 | rate_limit_regression: bool |
| 147 | output_stable: bool |
| 148 | partial_data: bool |
| 149 | baseline_complete: bool |
| 150 | shard_complete: bool |
| 151 | verdict: str |
| 152 | |
| 153 | @classmethod |
| 154 | def from_comparison(cls, comparison: dict[str, Any]) -> "ExperimentReport": |
| 155 | speedup_pct = comparison.get("speedup_pct") |
| 156 | api_growth_pct = comparison.get("api_growth_pct") |
| 157 | rate_limit_regression = bool( |
| 158 | comparison.get( |
| 159 | "rate_limit_regression", comparison.get("secondary_rate_limit_regression", False) |
| 160 | ) |
| 161 | ) |
| 162 | output_stable = bool(comparison.get("output_stable", False)) |
| 163 | partial_data = bool(comparison.get("partial_data", False)) |
| 164 | baseline_complete = bool(comparison.get("baseline_complete", True)) |
| 165 | shard_complete = bool(comparison.get("shard_complete", True)) |
| 166 | if partial_data or not baseline_complete or not shard_complete: |
| 167 | verdict = "inconclusive" |
| 168 | elif ( |
| 169 | speedup_pct is not None |
| 170 | and api_growth_pct is not None |
| 171 | and output_stable |
| 172 | and not rate_limit_regression |
| 173 | and float(speedup_pct) >= 25.0 |
| 174 | and float(api_growth_pct) <= 10.0 |
| 175 | ): |
| 176 | verdict = "pass" |
| 177 | else: |
| 178 | verdict = "fail" |
| 179 | return cls( |
| 180 | speedup_pct=speedup_pct, |
| 181 | api_growth_pct=api_growth_pct, |
| 182 | rate_limit_regression=rate_limit_regression, |
| 183 | output_stable=output_stable, |
| 184 | partial_data=partial_data, |
| 185 | baseline_complete=baseline_complete, |
| 186 | shard_complete=shard_complete, |
| 187 | verdict=verdict, |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | class SharedQuotaTracker: |
| 192 | def __init__(self, cap: int) -> None: |
| 193 | self.cap = max(0, int(cap)) |
| 194 | self.count = 0 |
| 195 | self._lock = threading.Lock() |
| 196 | |
| 197 | def increment(self) -> bool: |
| 198 | with self._lock: |
| 199 | if self.count >= self.cap: |
| 200 | return False |
| 201 | self.count += 1 |
| 202 | return True |
| 203 | |
| 204 | def reset(self) -> None: |
| 205 | with self._lock: |
| 206 | self.count = 0 |
| 207 | |
| 208 | |
| 209 | class WallClockBudget: |
| 210 | def __init__(self, budget_s: float) -> None: |
| 211 | self.budget_s = float(budget_s) |
| 212 | self.started_at = time.monotonic() |
| 213 | |
| 214 | def elapsed_s(self) -> float: |
| 215 | return time.monotonic() - self.started_at |
| 216 | |
| 217 | def is_exceeded(self) -> bool: |
| 218 | return self.elapsed_s() > self.budget_s |
| 219 | |
| 220 | |
| 221 | class SharedQuotaCoordinator: |
| 222 | def __init__(self, api_hard_cap: int | None) -> None: |
| 223 | self.api_hard_cap = api_hard_cap |
| 224 | self.total_api_calls = 0 |
| 225 | self.global_backoff_until = 0.0 |
| 226 | self.secondary_rate_limit_hit = False |
| 227 | self.abort_reason: str | None = None |
| 228 | self._guardrail_events: list[GuardrailEvent] = [] |
| 229 | self._lock = threading.Lock() |
| 230 | |
| 231 | def before_request(self, shard_name: str, query: str, deadline: float | None) -> None: |
| 232 | self._raise_if_aborted() |
| 233 | self._raise_if_budget_exceeded(shard_name, deadline) |
| 234 | while True: |
| 235 | with self._lock: |
| 236 | resume_at = self.global_backoff_until |
| 237 | now = time.monotonic() |
| 238 | if resume_at <= now: |
| 239 | break |
| 240 | time.sleep(min(resume_at - now, 1.0)) |
| 241 | self._raise_if_aborted() |
| 242 | self._raise_if_budget_exceeded(shard_name, deadline) |
| 243 | # Atomically reserve one API call slot to prevent concurrent over-cap |
| 244 | with self._lock: |
| 245 | if self.api_hard_cap is not None and self.total_api_calls >= self.api_hard_cap: |
| 246 | if self.abort_reason is None: |
| 247 | self.abort_reason = ( |
| 248 | f"API budget cap reached before {shard_name} requested {query} " |
| 249 | f"({self.total_api_calls}/{self.api_hard_cap})." |
| 250 | ) |
| 251 | self._guardrail_events.append( |
| 252 | GuardrailEvent( |
| 253 | kind="api_budget_cap", |
| 254 | shard=shard_name, |
| 255 | message=self.abort_reason, |
| 256 | at=iso_timestamp(utc_now()), |
| 257 | details={ |
| 258 | "api_calls_used": self.total_api_calls, |
| 259 | "api_hard_cap": self.api_hard_cap, |
| 260 | }, |
| 261 | ) |
| 262 | ) |
| 263 | raise ExperimentAbort(self.abort_reason) |
| 264 | # Reserve capacity for this request before releasing the lock |
| 265 | self.total_api_calls += 1 |
| 266 | |
| 267 | def register_api_calls(self, delta: int, shard_name: str) -> None: |
| 268 | # Subtract the 1 call already reserved in before_request() |
| 269 | additional = delta - 1 |
| 270 | if additional <= 0: |
| 271 | return |
| 272 | with self._lock: |
| 273 | self.total_api_calls += additional |
| 274 | if ( |
| 275 | self.api_hard_cap is not None |
| 276 | and self.total_api_calls > self.api_hard_cap |
| 277 | and self.abort_reason is None |
| 278 | ): |
| 279 | self.abort_reason = ( |
| 280 | f"API budget cap exceeded by {shard_name} " |
| 281 | f"({self.total_api_calls}/{self.api_hard_cap})." |
| 282 | ) |
| 283 | self._guardrail_events.append( |
| 284 | GuardrailEvent( |
| 285 | kind="api_budget_cap", |
| 286 | shard=shard_name, |
| 287 | message=self.abort_reason, |
| 288 | at=iso_timestamp(utc_now()), |
| 289 | details={ |
| 290 | "api_calls_used": self.total_api_calls, |
| 291 | "api_hard_cap": self.api_hard_cap, |
| 292 | }, |
| 293 | ) |
| 294 | ) |
| 295 | |
| 296 | def register_backoff( |
| 297 | self, shard_name: str, delay: float, reason: str, *, secondary: bool = False |
| 298 | ) -> None: |
| 299 | delay = max(delay, 1.0) |
| 300 | with self._lock: |
| 301 | self.global_backoff_until = max(self.global_backoff_until, time.monotonic() + delay) |
| 302 | self._guardrail_events.append( |
| 303 | GuardrailEvent( |
| 304 | kind="secondary_rate_limit" if secondary else "rate_limit_backoff", |
| 305 | shard=shard_name, |
| 306 | message=reason, |
| 307 | at=iso_timestamp(utc_now()), |
| 308 | details={"delay_seconds": round(delay, 2)}, |
| 309 | ) |
| 310 | ) |
| 311 | if secondary: |
| 312 | self.secondary_rate_limit_hit = True |
| 313 | self.abort_reason = reason |
| 314 | |
| 315 | def record_budget_exceeded(self, shard_name: str, budget_seconds: int) -> None: |
| 316 | with self._lock: |
| 317 | self._guardrail_events.append( |
| 318 | GuardrailEvent( |
| 319 | kind="wall_clock_budget", |
| 320 | shard=shard_name, |
| 321 | message=f"{shard_name} exhausted its {budget_seconds}s wall-clock budget.", |
| 322 | at=iso_timestamp(utc_now()), |
| 323 | details={"budget_seconds": budget_seconds}, |
| 324 | ) |
| 325 | ) |
| 326 | |
| 327 | def record_redistribution(self, shard_name: str, repo_count: int) -> None: |
| 328 | if repo_count <= 0: |
| 329 | return |
| 330 | with self._lock: |
| 331 | self._guardrail_events.append( |
| 332 | GuardrailEvent( |
| 333 | kind="work_redistributed", |
| 334 | shard=shard_name, |
| 335 | message=f"Redistributed {repo_count} remaining repositories from {shard_name}.", |
| 336 | at=iso_timestamp(utc_now()), |
| 337 | details={"repo_count": repo_count}, |
| 338 | ) |
| 339 | ) |
| 340 | |
| 341 | def guardrail_events(self) -> list[dict[str, Any]]: |
| 342 | with self._lock: |
| 343 | return [event.to_dict() for event in self._guardrail_events] |
| 344 | |
| 345 | def _raise_if_aborted(self) -> None: |
| 346 | with self._lock: |
| 347 | if self.abort_reason is not None: |
| 348 | raise ExperimentAbort(self.abort_reason) |
| 349 | |
| 350 | def _raise_if_budget_exceeded(self, shard_name: str, deadline: float | None) -> None: |
| 351 | if deadline is not None and time.monotonic() >= deadline: |
| 352 | raise ShardBudgetExceeded(f"{shard_name} exceeded its wall-clock budget.") |
| 353 | |
| 354 | |
| 355 | class InstrumentedGitHubClient(GitHubClient): |
| 356 | def __init__( |
| 357 | self, |
| 358 | token: str, |
| 359 | *, |
| 360 | cache_dir: Path, |
| 361 | shard_name: str, |
| 362 | coordinator: SharedQuotaCoordinator | None = None, |
| 363 | deadline: float | None = None, |
| 364 | timeout: int = 30, |
| 365 | max_retries: int = 6, |
| 366 | ) -> None: |
| 367 | self._api_calls_used = 0 |
| 368 | self.shard_name = shard_name |
| 369 | self.coordinator = coordinator |
| 370 | self.deadline = deadline |
| 371 | self.rate_limit_events = 0 |
| 372 | self.secondary_rate_limit_events = 0 |
| 373 | super().__init__(token, cache_dir=cache_dir, timeout=timeout, max_retries=max_retries) |
| 374 | |
| 375 | @property |
| 376 | def api_calls_used(self) -> int: |
| 377 | return self._api_calls_used |
| 378 | |
| 379 | @api_calls_used.setter |
| 380 | def api_calls_used(self, value: int) -> None: |
| 381 | prior = getattr(self, "_api_calls_used", 0) |
| 382 | self._api_calls_used = value |
| 383 | delta = value - prior |
| 384 | if self.coordinator is not None and delta > 0: |
| 385 | self.coordinator.register_api_calls(delta, self.shard_name) |
| 386 | |
| 387 | def _pause_for_rate_limit(self, query: str) -> None: |
| 388 | if self.coordinator is not None: |
| 389 | self.coordinator.before_request(self.shard_name, query, self.deadline) |
| 390 | elif self.deadline is not None and time.monotonic() >= self.deadline: |
| 391 | raise ShardBudgetExceeded(f"{self.shard_name} exceeded its wall-clock budget.") |
| 392 | super()._pause_for_rate_limit(query) |
| 393 | |
| 394 | def _respect_min_interval(self, url: str) -> None: |
| 395 | if self.deadline is not None and time.monotonic() >= self.deadline: |
| 396 | raise ShardBudgetExceeded(f"{self.shard_name} exceeded its wall-clock budget.") |
| 397 | super()._respect_min_interval(url) |
| 398 | |
| 399 | def _sleep_before_retry( |
| 400 | self, |
| 401 | attempt: int, |
| 402 | headers: dict[str, str] | None, |
| 403 | body: str, |
| 404 | query: str, |
| 405 | retry_limit: int, |
| 406 | max_delay_seconds: float, |
| 407 | ) -> None: |
| 408 | lowered = body.lower() |
| 409 | retry_after = None |
| 410 | if headers and headers.get("Retry-After"): |
| 411 | try: |
| 412 | retry_after = max(float(headers["Retry-After"]), 1.0) |
| 413 | except ValueError: |
| 414 | retry_after = None |
| 415 | reset_delay = self._reset_delay(headers) |
| 416 | delay = retry_after or reset_delay or min(2**attempt, max_delay_seconds) |
| 417 | if "secondary rate limit" in lowered: |
| 418 | self.rate_limit_events += 1 |
| 419 | self.secondary_rate_limit_events += 1 |
| 420 | reason = f"{self.shard_name} hit a secondary rate limit while requesting {query}." |
| 421 | if self.coordinator is not None: |
| 422 | self.coordinator.register_backoff( |
| 423 | self.shard_name, max(delay, 8.0), reason, secondary=True |
| 424 | ) |
| 425 | raise ExperimentAbort(reason) |
| 426 | if headers and ( |
| 427 | headers.get("Retry-After") is not None |
| 428 | or headers.get("X-RateLimit-Remaining") == "0" |
| 429 | or (self.rate_limit_remaining is not None and self.rate_limit_remaining <= 0) |
| 430 | ): |
| 431 | self.rate_limit_events += 1 |
| 432 | if self.coordinator is not None: |
| 433 | self.coordinator.register_backoff( |
| 434 | self.shard_name, |
| 435 | max(delay, 1.0), |
| 436 | f"{self.shard_name} backing off before retrying {query}.", |
| 437 | ) |
| 438 | super()._sleep_before_retry(attempt, headers, body, query, retry_limit, max_delay_seconds) |
| 439 | |
| 440 | |
| 441 | def parse_args() -> argparse.Namespace: |
| 442 | parser = argparse.ArgumentParser(description=__doc__) |
| 443 | parser.add_argument("--since", required=True, help="UTC crawl window start date (YYYY-MM-DD).") |
| 444 | parser.add_argument("--as-of", required=True, help="UTC crawl window end date (YYYY-MM-DD).") |
| 445 | parser.add_argument( |
| 446 | "--max-results", type=int, default=250, help="Maximum repositories per query." |
| 447 | ) |
| 448 | parser.add_argument("--topic", default=None, help="Optional topic id.") |
| 449 | parser.add_argument("--config", default=None, help="Optional crawl topic config file.") |
| 450 | parser.add_argument( |
| 451 | "--shards", type=int, default=3, help="Total shards including search shards." |
| 452 | ) |
| 453 | parser.add_argument("--wall-clock-budget", type=int, default=DEFAULT_WALL_CLOCK_BUDGET) |
| 454 | parser.add_argument( |
| 455 | "--api-budget-multiplier", type=float, default=DEFAULT_API_BUDGET_MULTIPLIER |
| 456 | ) |
| 457 | parser.add_argument("--output-dir", default=str(EXPERIMENT_ROOT)) |
| 458 | parser.add_argument("--experiment-id", default=None) |
| 459 | return parser.parse_args() |
| 460 | |
| 461 | |
| 462 | def build_context(args: argparse.Namespace, experiment_dir: Path) -> CrawlContext: |
| 463 | crawled_at = utc_now() |
| 464 | since = datetime.strptime(args.since, "%Y-%m-%d").replace(tzinfo=UTC) |
| 465 | window_end = datetime.strptime(args.as_of, "%Y-%m-%d").replace(tzinfo=UTC) |
| 466 | topic_id = args.topic |
| 467 | max_results = max(1, min(int(args.max_results), 1000)) |
| 468 | config_args = argparse.Namespace( |
| 469 | since=args.since, |
| 470 | as_of=args.as_of, |
| 471 | max_results=max_results, |
| 472 | output=str(experiment_dir / "baseline-raw.json"), |
| 473 | topic=topic_id, |
| 474 | config=args.config, |
| 475 | ) |
| 476 | return CrawlContext( |
| 477 | args=args, |
| 478 | topic_id=topic_id, |
| 479 | topic_raw=raw_dir(topic_id), |
| 480 | topic_snapshots=snapshots_dir(topic_id), |
| 481 | topic_cache=cache_dir(topic_id), |
| 482 | crawled_at=crawled_at, |
| 483 | run_started_at=crawled_at, |
| 484 | since=since, |
| 485 | window_end=window_end, |
| 486 | week=week_slug(window_end), |
| 487 | max_results=max_results, |
| 488 | config_checksum=github_crawl_config_checksum(config_args, since, window_end, max_results), |
| 489 | current_code_sha=os.environ.get("CRAWLER_CODE_SHA", ""), |
| 490 | source_refresh_policy="force-refresh", |
| 491 | baseline_output_path=experiment_dir / "baseline-raw.json", |
| 492 | baseline_snapshot_path=experiment_dir / "baseline-stars.json", |
| 493 | shard_output_path=experiment_dir / "shard-raw.json", |
| 494 | shard_snapshot_path=experiment_dir / "shard-stars.json", |
| 495 | ) |
| 496 | |
| 497 | |
| 498 | def next_experiment_id(output_dir: Path) -> str: |
| 499 | index = 0 |
| 500 | for candidate in output_dir.glob("shard-435-run-*"): |
| 501 | suffix = candidate.name.removeprefix("shard-435-run-") |
| 502 | if suffix.isdigit(): |
| 503 | index = max(index, int(suffix)) |
| 504 | return f"shard-435-run-{index + 1:03d}" |
| 505 | |
| 506 | |
| 507 | def build_search_plans(context: CrawlContext) -> list[SearchPlan]: |
| 508 | if context.args.config: |
| 509 | queries = load_topic_queries( |
| 510 | context.args.config, |
| 511 | { |
| 512 | "last_week": context.since.date().isoformat(), |
| 513 | "today": context.window_end.date().isoformat(), |
| 514 | }, |
| 515 | ) |
| 516 | return [ |
| 517 | SearchPlan( |
| 518 | shard_name="new-search", |
| 519 | repo_group="new", |
| 520 | primary_queries=list(queries["primary"]), |
| 521 | secondary_queries=list(queries["secondary"]), |
| 522 | min_repos_per_week=int(queries["min_repos_per_week"]), |
| 523 | ) |
| 524 | ] |
| 525 | return [ |
| 526 | SearchPlan( |
| 527 | shard_name="new-search", |
| 528 | repo_group="new", |
| 529 | primary_queries=[ |
| 530 | f"created:{context.since.date().isoformat()}..{context.window_end.date().isoformat()} stars:>50" |
| 531 | ], |
| 532 | ), |
| 533 | SearchPlan( |
| 534 | shard_name="trending-search", |
| 535 | repo_group="trending", |
| 536 | primary_queries=[ |
| 537 | f"pushed:{context.since.date().isoformat()}..{context.window_end.date().isoformat()} stars:>50" |
| 538 | ], |
| 539 | ), |
| 540 | ] |
| 541 | |
| 542 | |
| 543 | def run_search_plan( |
| 544 | plan: SearchPlan, |
| 545 | client: InstrumentedGitHubClient, |
| 546 | *, |
| 547 | max_results: int, |
| 548 | wall_clock_budget: int, |
| 549 | coordinator: SharedQuotaCoordinator | None, |
| 550 | ) -> tuple[list[dict[str, Any]], dict[str, Any]]: |
| 551 | candidates: list[dict[str, Any]] = [] |
| 552 | errors: list[str] = [] |
| 553 | completed = True |
| 554 | try: |
| 555 | for query in plan.primary_queries: |
| 556 | candidates.extend(client.search_repositories(query, max_results=max_results)) |
| 557 | if plan.secondary_queries and len(candidates) < plan.min_repos_per_week: |
| 558 | for query in plan.secondary_queries: |
| 559 | candidates.extend(client.search_repositories(query, max_results=max_results)) |
| 560 | except ShardBudgetExceeded as exc: |
| 561 | completed = False |
| 562 | errors.append(str(exc)) |
| 563 | if coordinator is not None: |
| 564 | coordinator.record_budget_exceeded(plan.shard_name, wall_clock_budget) |
| 565 | except ExperimentAbort: |
| 566 | raise |
| 567 | return candidates, { |
| 568 | "completed": completed, |
| 569 | "errors": errors + list(client.errors), |
| 570 | "api_calls": client.api_calls_used, |
| 571 | "cache_hits": client.cache_hits, |
| 572 | "stale_cache_hits": client.stale_cache_hits, |
| 573 | "rate_limit_events": client.rate_limit_events, |
| 574 | "secondary_rate_limit_events": client.secondary_rate_limit_events, |
| 575 | } |
| 576 | |
| 577 | |
| 578 | def chunk_validation_items( |
| 579 | items: list[ValidationItem], worker_count: int |
| 580 | ) -> list[list[ValidationItem]]: |
| 581 | if not items: |
| 582 | return [] |
| 583 | chunk_size = max(1, math.ceil(len(items) / max(worker_count * 2, 1))) |
| 584 | return [items[index : index + chunk_size] for index in range(0, len(items), chunk_size)] |
| 585 | |
| 586 | |
| 587 | def prepare_validation_items( |
| 588 | new_candidates: list[dict[str, Any]], |
| 589 | trending_candidates: list[dict[str, Any]], |
| 590 | worker_count: int, |
| 591 | ) -> tuple[Queue[list[ValidationItem]], dict[str, int]]: |
| 592 | duplicate_counts = {"new": 0, "trending": 0} |
| 593 | grouped_unique_items: list[ValidationItem] = [] |
| 594 | for repo_group, candidates in (("new", new_candidates), ("trending", trending_candidates)): |
| 595 | seen: set[str] = set() |
| 596 | for sequence, repo in enumerate(candidates): |
| 597 | full_name = repo.get("full_name") |
| 598 | if not full_name: |
| 599 | grouped_unique_items.append( |
| 600 | ValidationItem(repo_group=repo_group, sequence=sequence, repo=repo) |
| 601 | ) |
| 602 | continue |
| 603 | if full_name in seen: |
| 604 | duplicate_counts[repo_group] += 1 |
| 605 | continue |
| 606 | seen.add(full_name) |
| 607 | grouped_unique_items.append( |
| 608 | ValidationItem(repo_group=repo_group, sequence=sequence, repo=repo) |
| 609 | ) |
| 610 | queue: Queue[list[ValidationItem]] = Queue() |
| 611 | for chunk in chunk_validation_items(grouped_unique_items, worker_count): |
| 612 | queue.put(chunk) |
| 613 | return queue, duplicate_counts |
| 614 | |
| 615 | |
| 616 | def trending_stars_gained( |
| 617 | repo: dict[str, Any], |
| 618 | *, |
| 619 | previous_stars: dict[str, int] | None, |
| 620 | trending_cutoff: datetime | None, |
| 621 | ) -> int | None: |
| 622 | if previous_stars is None: |
| 623 | return None |
| 624 | previous_value = previous_stars.get(str(repo.get("full_name") or "")) |
| 625 | current_stars = int(repo.get("stargazers_count") or 0) |
| 626 | if previous_value is not None: |
| 627 | return max(current_stars - previous_value, 0) |
| 628 | if previous_stars and trending_cutoff is not None and isinstance(repo.get("created_at"), str): |
| 629 | created = datetime.fromisoformat(str(repo["created_at"]).replace("Z", "+00:00")) |
| 630 | if created >= trending_cutoff: |
| 631 | return current_stars |
| 632 | return None |
| 633 | |
| 634 | |
| 635 | def validate_item( |
| 636 | client: InstrumentedGitHubClient, |
| 637 | item: ValidationItem, |
| 638 | *, |
| 639 | previous_stars: dict[str, int] | None, |
| 640 | trending_cutoff: datetime | None, |
| 641 | ) -> tuple[ValidatedRecord | None, str | None]: |
| 642 | full_name = item.repo.get("full_name") |
| 643 | if not full_name: |
| 644 | return None, "missing_full_name" |
| 645 | skip_reason = significance_skip_reason(item.repo) |
| 646 | if skip_reason: |
| 647 | return None, skip_reason |
| 648 | try: |
| 649 | if not client.has_readme(full_name): |
| 650 | return None, "missing_readme" |
| 651 | except RuntimeError as exc: |
| 652 | client.record_error(f"README lookup failed for {full_name}: {exc}") |
| 653 | return None, "readme_lookup_failed" |
| 654 | stars_gained = None |
| 655 | if item.repo_group == "trending": |
| 656 | stars_gained = trending_stars_gained( |
| 657 | item.repo, |
| 658 | previous_stars=previous_stars, |
| 659 | trending_cutoff=trending_cutoff, |
| 660 | ) |
| 661 | return ValidatedRecord( |
| 662 | item.repo_group, item.sequence, to_repo_record(item.repo, stars_gained=stars_gained) |
| 663 | ), None |
| 664 | |
| 665 | |
| 666 | def sort_validated_records( |
| 667 | records: list[ValidatedRecord], |
| 668 | *, |
| 669 | previous_stars: dict[str, int] | None, |
| 670 | ) -> list[dict[str, Any]]: |
| 671 | if records and records[0].repo_group == "trending" and previous_stars: |
| 672 | ordered = sorted( |
| 673 | records, |
| 674 | key=lambda item: ( |
| 675 | -int( |
| 676 | item.record.get("stars_gained", -1) |
| 677 | if item.record.get("stars_gained") is not None |
| 678 | else -1 |
| 679 | ), |
| 680 | -int(item.record.get("stars") or 0), |
| 681 | item.sequence, |
| 682 | ), |
| 683 | ) |
| 684 | else: |
| 685 | ordered = sorted( |
| 686 | records, key=lambda item: (-int(item.record.get("stars") or 0), item.sequence) |
| 687 | ) |
| 688 | return [item.record for item in ordered] |
| 689 | |
| 690 | |
| 691 | def validation_worker( |
| 692 | worker_index: int, |
| 693 | queue: Queue[list[ValidationItem]], |
| 694 | client: InstrumentedGitHubClient, |
| 695 | *, |
| 696 | previous_stars: dict[str, int] | None, |
| 697 | trending_cutoff: datetime | None, |
| 698 | wall_clock_budget: int, |
| 699 | coordinator: SharedQuotaCoordinator, |
| 700 | ) -> dict[str, Any]: |
| 701 | shard_name = f"validate-{worker_index}" |
| 702 | results = {"new": [], "trending": []} |
| 703 | filters: dict[str, Counter[str]] = {"new": Counter(), "trending": Counter()} |
| 704 | completed = True |
| 705 | while True: |
| 706 | try: |
| 707 | chunk = queue.get_nowait() |
| 708 | except Empty: |
| 709 | break |
| 710 | current_index = 0 |
| 711 | try: |
| 712 | for current_index, item in enumerate(chunk): |
| 713 | validated, skip_reason = validate_item( |
| 714 | client, |
| 715 | item, |
| 716 | previous_stars=previous_stars, |
| 717 | trending_cutoff=trending_cutoff, |
| 718 | ) |
| 719 | if validated is not None: |
| 720 | results[item.repo_group].append(validated) |
| 721 | elif skip_reason is not None: |
| 722 | filters[item.repo_group][skip_reason] += 1 |
| 723 | except ShardBudgetExceeded: |
| 724 | completed = False |
| 725 | coordinator.record_budget_exceeded(shard_name, wall_clock_budget) |
| 726 | remaining_items = chunk[current_index:] |
| 727 | if remaining_items: |
| 728 | queue.put(remaining_items) |
| 729 | coordinator.record_redistribution(shard_name, len(remaining_items)) |
| 730 | break |
| 731 | except ExperimentAbort: |
| 732 | raise |
| 733 | return { |
| 734 | "completed": completed, |
| 735 | "results": results, |
| 736 | "filters": {name: dict(counter) for name, counter in filters.items()}, |
| 737 | "errors": list(client.errors), |
| 738 | "api_calls": client.api_calls_used, |
| 739 | "cache_hits": client.cache_hits, |
| 740 | "stale_cache_hits": client.stale_cache_hits, |
| 741 | "rate_limit_events": client.rate_limit_events, |
| 742 | "secondary_rate_limit_events": client.secondary_rate_limit_events, |
| 743 | } |
| 744 | |
| 745 | |
| 746 | def build_payload( |
| 747 | *, |
| 748 | context: CrawlContext, |
| 749 | output_path: Path, |
| 750 | snapshot_path: Path, |
| 751 | new_repos: list[dict[str, Any]], |
| 752 | trending_repos: list[dict[str, Any]], |
| 753 | new_candidates: list[dict[str, Any]], |
| 754 | trending_candidates: list[dict[str, Any]], |
| 755 | api_calls: int, |
| 756 | cache_hits: int, |
| 757 | stale_cache_hits: int, |
| 758 | rate_limit_limit: int | None, |
| 759 | rate_limit_remaining: int | None, |
| 760 | rate_limit_reset: int | None, |
| 761 | rate_limit_resource: str | None, |
| 762 | partial_failures: list[str], |
| 763 | filter_summary: dict[str, dict[str, int]], |
| 764 | run_mode: str, |
| 765 | ) -> tuple[dict[str, Any], dict[str, Any]]: |
| 766 | star_snapshot = build_star_snapshot(new_candidates, trending_candidates) |
| 767 | snapshot_payload = { |
| 768 | "week": context.week, |
| 769 | "captured_at": iso_timestamp(context.crawled_at), |
| 770 | "repository_count": len(star_snapshot), |
| 771 | "stars": star_snapshot, |
| 772 | } |
| 773 | payload = { |
| 774 | "week": context.week, |
| 775 | "crawled_at": iso_timestamp(context.crawled_at), |
| 776 | "new_repos": new_repos, |
| 777 | "trending_repos": trending_repos, |
| 778 | "signals": build_signals(new_repos, trending_repos), |
| 779 | "metadata": { |
| 780 | "api_calls_used": api_calls, |
| 781 | "cache_hits": cache_hits, |
| 782 | "stale_cache_hits": stale_cache_hits, |
| 783 | "rate_limit_limit": rate_limit_limit, |
| 784 | "rate_limit_remaining": rate_limit_remaining, |
| 785 | "rate_limit_reset": rate_limit_reset, |
| 786 | "rate_limit_resource": rate_limit_resource, |
| 787 | "partial_failures": partial_failures, |
| 788 | "run_id": f"local-{run_mode}", |
| 789 | "crawl_window": { |
| 790 | "since": context.since.date().isoformat(), |
| 791 | "until": context.window_end.date().isoformat(), |
| 792 | }, |
| 793 | "crawl_config_checksum": context.config_checksum, |
| 794 | "schema_checksum": github_schema_checksum(), |
| 795 | "same_day_reuse": { |
| 796 | "status": "not_reused", |
| 797 | "source": "github", |
| 798 | "source_id": "github-search", |
| 799 | }, |
| 800 | "filter_summary": filter_summary, |
| 801 | "snapshot_path": snapshot_path.as_posix(), |
| 802 | "source_refresh_policy": context.source_refresh_policy, |
| 803 | "crawler_code_sha": context.current_code_sha, |
| 804 | "experiment_mode": run_mode, |
| 805 | "output_path": output_path.as_posix(), |
| 806 | }, |
| 807 | } |
| 808 | payload["metadata"]["artifact_checksum"] = github_artifact_checksum(payload) |
| 809 | validate_payload(payload) |
| 810 | return payload, snapshot_payload |
| 811 | |
| 812 | |
| 813 | def run_baseline(context: CrawlContext, token: str) -> RunResult: |
| 814 | started_at = time.monotonic() |
| 815 | search_plans = build_search_plans(context) |
| 816 | client = InstrumentedGitHubClient(token, cache_dir=context.topic_cache, shard_name="baseline") |
| 817 | previous_stars = load_previous_star_snapshot( |
| 818 | context.topic_snapshots, |
| 819 | context.week, |
| 820 | context.baseline_output_path.parent, |
| 821 | context.topic_raw, |
| 822 | RAW_ROOT, |
| 823 | ) |
| 824 | new_candidates: list[dict[str, Any]] = [] |
| 825 | trending_candidates: list[dict[str, Any]] = [] |
| 826 | if context.args.config: |
| 827 | new_candidates, _ = run_search_plan( |
| 828 | search_plans[0], |
| 829 | client, |
| 830 | max_results=context.max_results, |
| 831 | wall_clock_budget=context.args.wall_clock_budget, |
| 832 | coordinator=None, |
| 833 | ) |
| 834 | else: |
| 835 | for plan in search_plans: |
| 836 | candidates, _ = run_search_plan( |
| 837 | plan, |
| 838 | client, |
| 839 | max_results=context.max_results, |
| 840 | wall_clock_budget=context.args.wall_clock_budget, |
| 841 | coordinator=None, |
| 842 | ) |
| 843 | if plan.repo_group == "new": |
| 844 | new_candidates = candidates |
| 845 | else: |
| 846 | trending_candidates = candidates |
| 847 | new_repos, new_filters = collect_repositories(client, new_candidates) |
| 848 | trending_repos, trending_filters = collect_repositories( |
| 849 | client, |
| 850 | trending_candidates, |
| 851 | previous_stars=previous_stars, |
| 852 | trending_cutoff=context.since, |
| 853 | ) |
| 854 | payload, snapshot_payload = build_payload( |
| 855 | context=context, |
| 856 | output_path=context.baseline_output_path, |
| 857 | snapshot_path=context.baseline_snapshot_path, |
| 858 | new_repos=new_repos, |
| 859 | trending_repos=trending_repos, |
| 860 | new_candidates=new_candidates, |
| 861 | trending_candidates=trending_candidates, |
| 862 | api_calls=client.api_calls_used, |
| 863 | cache_hits=client.cache_hits, |
| 864 | stale_cache_hits=client.stale_cache_hits, |
| 865 | rate_limit_limit=client.rate_limit_limit, |
| 866 | rate_limit_remaining=client.rate_limit_remaining, |
| 867 | rate_limit_reset=client.rate_limit_reset, |
| 868 | rate_limit_resource=client.rate_limit_resource, |
| 869 | partial_failures=list(client.errors), |
| 870 | filter_summary={"new_repos": new_filters, "trending_repos": trending_filters}, |
| 871 | run_mode="baseline", |
| 872 | ) |
| 873 | write_payload(context.baseline_output_path, payload) |
| 874 | write_payload(context.baseline_snapshot_path, snapshot_payload) |
| 875 | return RunResult( |
| 876 | name="baseline", |
| 877 | payload=payload, |
| 878 | snapshot_payload=snapshot_payload, |
| 879 | api_calls=client.api_calls_used, |
| 880 | cache_hits=client.cache_hits, |
| 881 | stale_cache_hits=client.stale_cache_hits, |
| 882 | rate_limit_events=client.rate_limit_events, |
| 883 | secondary_rate_limit_events=client.secondary_rate_limit_events, |
| 884 | partial_failures=list(client.errors), |
| 885 | wall_clock_s=round(time.monotonic() - started_at, 3), |
| 886 | shards_used=1, |
| 887 | completed=not bool(client.errors), |
| 888 | guardrail_events=[], |
| 889 | ) |
| 890 | |
| 891 | |
| 892 | def run_sharded(context: CrawlContext, token: str, baseline_api_calls: int) -> RunResult: |
| 893 | started_at = time.monotonic() |
| 894 | search_plans = build_search_plans(context) |
| 895 | validation_workers = max(1, int(context.args.shards) - len(search_plans)) |
| 896 | api_hard_cap = max( |
| 897 | 1, math.floor(baseline_api_calls * float(context.args.api_budget_multiplier)) |
| 898 | ) |
| 899 | coordinator = SharedQuotaCoordinator(api_hard_cap) |
| 900 | previous_stars = load_previous_star_snapshot( |
| 901 | context.topic_snapshots, |
| 902 | context.week, |
| 903 | context.shard_output_path.parent, |
| 904 | context.topic_raw, |
| 905 | RAW_ROOT, |
| 906 | ) |
| 907 | aggregated_errors: list[str] = [] |
| 908 | aggregated_api_calls = 0 |
| 909 | aggregated_cache_hits = 0 |
| 910 | aggregated_stale_cache_hits = 0 |
| 911 | aggregated_rate_limit_events = 0 |
| 912 | aggregated_secondary_rate_limit_events = 0 |
| 913 | search_results: dict[str, list[dict[str, Any]]] = {"new": [], "trending": []} |
| 914 | |
| 915 | try: |
| 916 | with ThreadPoolExecutor(max_workers=max(len(search_plans), 1)) as pool: |
| 917 | futures = { |
| 918 | pool.submit( |
| 919 | run_search_plan, |
| 920 | plan, |
| 921 | InstrumentedGitHubClient( |
| 922 | token, |
| 923 | cache_dir=context.topic_cache, |
| 924 | shard_name=plan.shard_name, |
| 925 | coordinator=coordinator, |
| 926 | deadline=time.monotonic() + int(context.args.wall_clock_budget), |
| 927 | ), |
| 928 | max_results=context.max_results, |
| 929 | wall_clock_budget=int(context.args.wall_clock_budget), |
| 930 | coordinator=coordinator, |
| 931 | ): plan |
| 932 | for plan in search_plans |
| 933 | } |
| 934 | for future in as_completed(futures): |
| 935 | plan = futures[future] |
| 936 | candidates, metrics = future.result() |
| 937 | search_results[plan.repo_group] = candidates |
| 938 | aggregated_errors.extend(metrics["errors"]) |
| 939 | aggregated_api_calls += metrics["api_calls"] |
| 940 | aggregated_cache_hits += metrics["cache_hits"] |
| 941 | aggregated_stale_cache_hits += metrics["stale_cache_hits"] |
| 942 | aggregated_rate_limit_events += metrics["rate_limit_events"] |
| 943 | aggregated_secondary_rate_limit_events += metrics["secondary_rate_limit_events"] |
| 944 | except ExperimentAbort as exc: |
| 945 | aggregated_errors.append(str(exc)) |
| 946 | |
| 947 | validation_queue, duplicate_counts = prepare_validation_items( |
| 948 | search_results["new"], search_results["trending"], validation_workers |
| 949 | ) |
| 950 | validated_records: dict[str, list[ValidatedRecord]] = {"new": [], "trending": []} |
| 951 | filter_summary: dict[str, Counter[str]] = { |
| 952 | "new_repos": Counter({"duplicate": duplicate_counts["new"]}), |
| 953 | "trending_repos": Counter({"duplicate": duplicate_counts["trending"]}), |
| 954 | } |
| 955 | if coordinator.abort_reason is None: |
| 956 | try: |
| 957 | with ThreadPoolExecutor(max_workers=validation_workers) as pool: |
| 958 | futures = { |
| 959 | pool.submit( |
| 960 | validation_worker, |
| 961 | index + 1, |
| 962 | validation_queue, |
| 963 | InstrumentedGitHubClient( |
| 964 | token, |
| 965 | cache_dir=context.topic_cache, |
| 966 | shard_name=f"validate-{index + 1}", |
| 967 | coordinator=coordinator, |
| 968 | deadline=time.monotonic() + int(context.args.wall_clock_budget), |
| 969 | ), |
| 970 | previous_stars=previous_stars, |
| 971 | trending_cutoff=context.since, |
| 972 | wall_clock_budget=int(context.args.wall_clock_budget), |
| 973 | coordinator=coordinator, |
| 974 | ): index + 1 |
| 975 | for index in range(validation_workers) |
| 976 | } |
| 977 | for future in as_completed(futures): |
| 978 | result = future.result() |
| 979 | validated_records["new"].extend(result["results"]["new"]) |
| 980 | validated_records["trending"].extend(result["results"]["trending"]) |
| 981 | filter_summary["new_repos"].update(result["filters"]["new"]) |
| 982 | filter_summary["trending_repos"].update(result["filters"]["trending"]) |
| 983 | aggregated_errors.extend(result["errors"]) |
| 984 | aggregated_api_calls += result["api_calls"] |
| 985 | aggregated_cache_hits += result["cache_hits"] |
| 986 | aggregated_stale_cache_hits += result["stale_cache_hits"] |
| 987 | aggregated_rate_limit_events += result["rate_limit_events"] |
| 988 | aggregated_secondary_rate_limit_events += result["secondary_rate_limit_events"] |
| 989 | except ExperimentAbort as exc: |
| 990 | aggregated_errors.append(str(exc)) |
| 991 | |
| 992 | new_repos = sort_validated_records(validated_records["new"], previous_stars=None) |
| 993 | trending_repos = sort_validated_records( |
| 994 | validated_records["trending"], previous_stars=previous_stars |
| 995 | ) |
| 996 | payload, snapshot_payload = build_payload( |
| 997 | context=context, |
| 998 | output_path=context.shard_output_path, |
| 999 | snapshot_path=context.shard_snapshot_path, |
| 1000 | new_repos=new_repos, |
| 1001 | trending_repos=trending_repos, |
| 1002 | new_candidates=search_results["new"], |
| 1003 | trending_candidates=search_results["trending"], |
| 1004 | api_calls=aggregated_api_calls, |
| 1005 | cache_hits=aggregated_cache_hits, |
| 1006 | stale_cache_hits=aggregated_stale_cache_hits, |
| 1007 | rate_limit_limit=None, |
| 1008 | rate_limit_remaining=None, |
| 1009 | rate_limit_reset=None, |
| 1010 | rate_limit_resource="mixed", |
| 1011 | partial_failures=aggregated_errors, |
| 1012 | filter_summary={name: dict(counter) for name, counter in filter_summary.items()}, |
| 1013 | run_mode="shard", |
| 1014 | ) |
| 1015 | write_payload(context.shard_output_path, payload) |
| 1016 | write_payload(context.shard_snapshot_path, snapshot_payload) |
| 1017 | return RunResult( |
| 1018 | name="shard", |
| 1019 | payload=payload, |
| 1020 | snapshot_payload=snapshot_payload, |
| 1021 | api_calls=aggregated_api_calls, |
| 1022 | cache_hits=aggregated_cache_hits, |
| 1023 | stale_cache_hits=aggregated_stale_cache_hits, |
| 1024 | rate_limit_events=aggregated_rate_limit_events, |
| 1025 | secondary_rate_limit_events=aggregated_secondary_rate_limit_events, |
| 1026 | partial_failures=aggregated_errors, |
| 1027 | wall_clock_s=round(time.monotonic() - started_at, 3), |
| 1028 | shards_used=len(search_plans) + validation_workers, |
| 1029 | completed=not aggregated_errors |
| 1030 | and validation_queue.empty() |
| 1031 | and not coordinator.secondary_rate_limit_hit, |
| 1032 | guardrail_events=coordinator.guardrail_events(), |
| 1033 | ) |
| 1034 | |
| 1035 | |
| 1036 | def canonicalize_payload(payload: dict[str, Any]) -> bytes: |
| 1037 | canonical = { |
| 1038 | "week": payload.get("week"), |
| 1039 | "new_repos": payload.get("new_repos", []), |
| 1040 | "trending_repos": payload.get("trending_repos", []), |
| 1041 | "signals": payload.get("signals", {}), |
| 1042 | } |
| 1043 | return json.dumps(canonical, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode( |
| 1044 | "utf-8" |
| 1045 | ) |
| 1046 | |
| 1047 | |
| 1048 | def canonicalize_snapshot(snapshot_payload: dict[str, Any]) -> bytes: |
| 1049 | canonical = { |
| 1050 | "week": snapshot_payload.get("week"), |
| 1051 | "repository_count": snapshot_payload.get("repository_count"), |
| 1052 | "stars": snapshot_payload.get("stars", {}), |
| 1053 | } |
| 1054 | return json.dumps(canonical, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode( |
| 1055 | "utf-8" |
| 1056 | ) |
| 1057 | |
| 1058 | |
| 1059 | def _normalize_for_comparison(value: Any) -> Any: |
| 1060 | if isinstance(value, dict): |
| 1061 | return { |
| 1062 | key: _normalize_for_comparison(item) |
| 1063 | for key, item in sorted(value.items()) |
| 1064 | if key not in {"crawled_at", "captured_at", "report_generated_at"} |
| 1065 | } |
| 1066 | if isinstance(value, list): |
| 1067 | return [_normalize_for_comparison(item) for item in value] |
| 1068 | return value |
| 1069 | |
| 1070 | |
| 1071 | def deterministic_merge(shard_results: list[ShardResult]) -> dict[str, Any]: |
| 1072 | merged: dict[str, dict[str, Any]] = {} |
| 1073 | for result in shard_results: |
| 1074 | for repo in result.repos_found: |
| 1075 | full_name = str(repo.get("full_name") or "") |
| 1076 | if full_name and full_name not in merged: |
| 1077 | merged[full_name] = repo |
| 1078 | repos = sorted( |
| 1079 | merged.values(), |
| 1080 | key=lambda item: ( |
| 1081 | -(int(item.get("stars_gained", -1) or -1)), |
| 1082 | -(int(item.get("stars") or 0)), |
| 1083 | str(item.get("full_name") or ""), |
| 1084 | ), |
| 1085 | ) |
| 1086 | return {"repos": repos} |
| 1087 | |
| 1088 | |
| 1089 | def compare_results(baseline: dict[str, Any], shard: dict[str, Any]) -> dict[str, Any]: |
| 1090 | baseline_wall = float(baseline.get("wall_clock_s", baseline.get("elapsed_s", 0.0)) or 0.0001) |
| 1091 | shard_wall = float(shard.get("wall_clock_s", shard.get("elapsed_s", 0.0)) or 0.0) |
| 1092 | baseline_api = int(baseline.get("api_calls", baseline.get("api_calls_used", 0)) or 0) |
| 1093 | shard_api = int(shard.get("api_calls", shard.get("api_calls_used", 0)) or 0) |
| 1094 | baseline_output = _normalize_for_comparison( |
| 1095 | baseline.get("canonical_output", baseline.get("output", {})) |
| 1096 | ) |
| 1097 | shard_output = _normalize_for_comparison(shard.get("canonical_output", shard.get("output", {}))) |
| 1098 | return { |
| 1099 | "speedup_pct": round(((baseline_wall - shard_wall) / baseline_wall) * 100, 2), |
| 1100 | "api_growth_pct": round(((shard_api - baseline_api) / max(baseline_api, 1)) * 100, 2), |
| 1101 | "output_stable": baseline_output == shard_output, |
| 1102 | "rate_limit_regression": int(shard.get("rate_limit_events", 0) or 0) |
| 1103 | > int(baseline.get("rate_limit_events", 0) or 0), |
| 1104 | "partial_data": bool( |
| 1105 | shard.get("partial_data", False) or baseline.get("partial_data", False) |
| 1106 | ), |
| 1107 | "baseline_complete": not bool(baseline.get("partial_data", False)), |
| 1108 | "shard_complete": not bool(shard.get("partial_data", False)), |
| 1109 | } |
| 1110 | |
| 1111 | |
| 1112 | def build_report(experiment_id: str, baseline: RunResult, shard: RunResult) -> dict[str, Any]: |
| 1113 | baseline_wall = baseline.wall_clock_s or 0.0001 |
| 1114 | speedup_pct = round(((baseline_wall - shard.wall_clock_s) / baseline_wall) * 100, 2) |
| 1115 | api_growth_pct = round( |
| 1116 | ((shard.api_calls - baseline.api_calls) / max(baseline.api_calls, 1)) * 100, 2 |
| 1117 | ) |
| 1118 | output_stable = canonicalize_payload(baseline.payload) == canonicalize_payload( |
| 1119 | shard.payload |
| 1120 | ) and canonicalize_snapshot(baseline.snapshot_payload) == canonicalize_snapshot( |
| 1121 | shard.snapshot_payload |
| 1122 | ) |
| 1123 | baseline_incomplete = bool(baseline.partial_failures) or not baseline.completed |
| 1124 | shard_incomplete = bool(shard.partial_failures) or not shard.completed |
| 1125 | report_card = ExperimentReport.from_comparison( |
| 1126 | { |
| 1127 | "speedup_pct": speedup_pct, |
| 1128 | "api_growth_pct": api_growth_pct, |
| 1129 | "rate_limit_regression": shard.secondary_rate_limit_events |
| 1130 | > baseline.secondary_rate_limit_events, |
| 1131 | "output_stable": output_stable, |
| 1132 | "partial_data": baseline_incomplete or shard_incomplete, |
| 1133 | "baseline_complete": not baseline_incomplete, |
| 1134 | "shard_complete": not shard_incomplete, |
| 1135 | } |
| 1136 | ) |
| 1137 | verdict = report_card.verdict |
| 1138 | if report_card.partial_data and any( |
| 1139 | event["kind"] == "secondary_rate_limit" for event in shard.guardrail_events |
| 1140 | ): |
| 1141 | verdict = "fail" |
| 1142 | return { |
| 1143 | "experiment_id": experiment_id, |
| 1144 | "baseline": { |
| 1145 | "wall_clock_s": baseline.wall_clock_s, |
| 1146 | "api_calls": baseline.api_calls, |
| 1147 | "rate_limit_events": baseline.rate_limit_events, |
| 1148 | "repos_new": len(baseline.payload.get("new_repos", [])), |
| 1149 | "repos_trending": len(baseline.payload.get("trending_repos", [])), |
| 1150 | }, |
| 1151 | "shard": { |
| 1152 | "wall_clock_s": shard.wall_clock_s, |
| 1153 | "api_calls": shard.api_calls, |
| 1154 | "rate_limit_events": shard.rate_limit_events, |
| 1155 | "shards_used": shard.shards_used, |
| 1156 | "repos_new": len(shard.payload.get("new_repos", [])), |
| 1157 | "repos_trending": len(shard.payload.get("trending_repos", [])), |
| 1158 | }, |
| 1159 | "comparison": { |
| 1160 | "speedup_pct": speedup_pct, |
| 1161 | "api_growth_pct": api_growth_pct, |
| 1162 | "output_stable": output_stable, |
| 1163 | "secondary_rate_limit_regression": shard.secondary_rate_limit_events |
| 1164 | > baseline.secondary_rate_limit_events, |
| 1165 | }, |
| 1166 | "verdict": verdict, |
| 1167 | "guardrail_events": shard.guardrail_events, |
| 1168 | } |
| 1169 | |
| 1170 | |
| 1171 | def main() -> int: |
| 1172 | args = parse_args() |
| 1173 | if args.shards < 3: |
| 1174 | print("--shards must be >= 3.", file=os.sys.stderr) |
| 1175 | return 1 |
| 1176 | if args.wall_clock_budget <= 0: |
| 1177 | print("--wall-clock-budget must be positive.", file=os.sys.stderr) |
| 1178 | return 1 |
| 1179 | if args.api_budget_multiplier < 1.0: |
| 1180 | print("--api-budget-multiplier must be >= 1.0.", file=os.sys.stderr) |
| 1181 | return 1 |
| 1182 | token = os.environ.get("GITHUB_TOKEN") |
| 1183 | if not token: |
| 1184 | print("GITHUB_TOKEN is required", file=os.sys.stderr) |
| 1185 | return 1 |
| 1186 | output_dir = Path(args.output_dir) |
| 1187 | output_dir.mkdir(parents=True, exist_ok=True) |
| 1188 | experiment_id = args.experiment_id or next_experiment_id(output_dir) |
| 1189 | experiment_dir = output_dir / experiment_id |
| 1190 | experiment_dir.mkdir(parents=True, exist_ok=True) |
| 1191 | context = build_context(args, experiment_dir) |
| 1192 | baseline = run_baseline(context, token) |
| 1193 | shard = run_sharded(context, token, baseline.api_calls) |
| 1194 | report = build_report(experiment_id, baseline, shard) |
| 1195 | report_path = experiment_dir / "report.json" |
| 1196 | write_payload(report_path, report) |
| 1197 | print( |
| 1198 | json.dumps( |
| 1199 | { |
| 1200 | "experiment_id": experiment_id, |
| 1201 | "verdict": report["verdict"], |
| 1202 | "speedup_pct": report["comparison"]["speedup_pct"], |
| 1203 | "api_growth_pct": report["comparison"]["api_growth_pct"], |
| 1204 | "output_stable": report["comparison"]["output_stable"], |
| 1205 | "report_path": report_path.as_posix(), |
| 1206 | }, |
| 1207 | ensure_ascii=False, |
| 1208 | ) |
| 1209 | ) |
| 1210 | return 0 if report["verdict"] != "fail" else 1 |
| 1211 | |
| 1212 | |
| 1213 | if __name__ == "__main__": |
| 1214 | raise SystemExit(main()) |