SPB Git forge
spboucher.ai ↗

spb/company-atlas

Public
28commits 1branches 0releases
7.7 MBsize
maindefault branch
10 days agolast push
Python 66.3% TypeScript 22.7% JavaScript 8.6% HTML 1.4% CSS 0.7%

onboarding: rolling worker pool (no batch barrier), discovery rate setting, onboard --limit 0 = all pending

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Simon-Pierre Boucher committed 12 days ago (Sep 12, 2026) parent f09c483

3 changed files +69 −47

modified src/companyatlas/commands/crawl.py +1 −1
@@ -26,7 +26,7 @@ def _table(title: str, columns: list[str], rows: list[list[Any]]) -> None:
26 26
27 27 def register(app: typer.Typer) -> None:
28 28 @app.command()
29 − def onboard(limit: int = 50, company: Annotated[str | None, typer.Option(help="slug / id / domain of one company")] = None,
29 + def onboard(limit: Annotated[int, typer.Option(help="max companies this run (0 = all pending)")] = 0, company: Annotated[str | None, typer.Option(help="slug / id / domain of one company")] = None,
30 30 concurrency: int | None = None, fetch_now: bool = False, dry_run: bool = False) -> None:
31 31 """Discover surfaces and create sensors for pending companies (or one company)."""
32 32 from companyatlas.services.discovery import onboard_pending
modified src/companyatlas/config.py +1 −0
@@ -50,6 +50,7 @@ class Settings(BaseSettings):
50 50 discovery_min_confidence: float = Field(0.55, alias="CA_DISCOVERY_MIN_CONFIDENCE")
51 51 discovery_max_sensors_per_company: int = Field(40, alias="CA_DISCOVERY_MAX_SENSORS")
52 52 onboarding_concurrency: int = Field(12, alias="CA_ONBOARDING_CONCURRENCY")
53 + discovery_rate_per_min: int = Field(40, alias="CA_DISCOVERY_RATE_PER_MIN") # same-domain spacing during discovery (1.5 s)
53 54
54 55 # ---------------------------------------------------------------- scheduling (seconds)
55 56 scheduler_tick_s: int = Field(15, alias="CA_SCHEDULER_TICK_S")
modified src/companyatlas/services/discovery.py +67 −46
@@ -152,7 +152,7 @@ async def _get(fetcher: Fetcher, url: str, *, max_bytes: int = PROBE_MAX_BYTES,
152 152 accept: str | None = None) -> FetchResult | None:
153 153 res.requests += 1
154 154 try:
155 − return await fetcher.get(url, max_bytes=max_bytes, respect_robots=respect_robots, retries=0, accept=accept)
155 + return await fetcher.get(url, max_bytes=max_bytes, respect_robots=respect_robots, retries=0, accept=accept, rate_per_min=settings.discovery_rate_per_min)
156 156 except NotModified:
157 157 return None
158 158 except (FetchError, BlockedError) as exc:
@@ -259,7 +259,7 @@ async def _phase_homepage(fetcher: Fetcher, company: dict[str, Any], res: Discov
259 259 break
260 260 res.requests += 1
261 261 try:
262 − r = await fetcher.get(url, max_bytes=HOMEPAGE_MAX_BYTES, retries=1)
262 + r = await fetcher.get(url, max_bytes=HOMEPAGE_MAX_BYTES, retries=1, rate_per_min=settings.discovery_rate_per_min)
263 263 except BlockedError as exc:
264 264 res.notes.append(f"homepage blocked: {exc.failure}")
265 265 res.error = f"{exc.failure}: {exc}"[:300]
@@ -689,59 +689,80 @@ async def finish_discover_job(job_id: str, *, ok: bool, error: str | None, attem
689 689 run_at = now() + make_interval(mins => :mins) where id = :id""", id=job_id, e=(error or "")[:500], mins=30 * attempts)
690 690
691 691
692 −async def onboard_pending(limit: int = 50, concurrency: int | None = None, *, fetcher: Fetcher | None = None, company_slug: str | None = None,
692 +async def onboard_pending(limit: int = 0, concurrency: int | None = None, *, fetcher: Fetcher | None = None, company_slug: str | None = None,
693 693 fetch_now: bool = False, dry_run: bool = False, worker: str = "onboard") -> dict[str, Any]:
694 − """Discover pending companies: queue jobs first (SKIP LOCKED), then `companies.onboarding_status='pending'` without a job."""
694 + """Discover pending companies with a rolling pool of `concurrency` workers (no batch barrier): each worker claims ONE target at a
695 + time — a `discover` queue job first (SKIP LOCKED), else a `companies.onboarding_status='pending'` row — until `limit` companies
696 + have been processed (`limit <= 0` = everything pending) or nothing is left. Safe to run from several processes/machines."""
695 697 concurrency = concurrency or settings.onboarding_concurrency
696 698 own_fetcher = fetcher is None
697 699 fetcher = fetcher or Fetcher()
698 700 if own_fetcher:
699 701 await fetcher.open()
700 702 stats = {"claimed": 0, "active": 0, "failed": 0, "no_website": 0, "sensors": 0}
701 − try:
702 − targets: list[tuple[dict[str, Any], dict[str, Any] | None]] = []
703 + lock = asyncio.Lock()
704 + budget = {"left": limit if limit and limit > 0 else float("inf")}
705 +
706 + async def process(comp: dict[str, Any], job: dict[str, Any] | None) -> None:
707 + res = await discover_company(comp, fetcher=fetcher, dry_run=dry_run, fetch_now=fetch_now)
708 + async with lock:
709 + stats[str(res.status)] = stats.get(str(res.status), 0) + 1
710 + stats["sensors"] += len(res.sensors)
711 + if job is not None:
712 + await finish_discover_job(job["id"], ok=res.status == OnboardingStatus.ACTIVE, error=res.error, attempts=job["attempts"], max_attempts=job["max_attempts"])
713 + log.info("company discovered", extra={"company": comp.get("slug"), "status": str(res.status), "sensors": len(res.sensors), "requests": res.requests,
714 + "ms": res.duration_ms, "error": res.error})
715 + if fetch_now and not dry_run and res.sensors:
716 + from companyatlas.services.pipeline import run_sensor_ids
717 +
718 + await run_sensor_ids([s["id"] for s in res.sensors], fetcher=fetcher, worker=worker)
719 +
720 + async def claim_one() -> tuple[dict[str, Any], dict[str, Any] | None] | None:
721 + async with lock:
722 + if budget["left"] <= 0:
723 + return None
724 + budget["left"] -= 1
725 + for j in await claim_discover_jobs(1, worker):
726 + cid = (j.get("payload") or {}).get("company_id") or j["key"].removeprefix("discover:")
727 + async with transaction() as conn:
728 + comp = await fetch_one(conn, "select * from companies where id = :id or slug = :id", id=cid)
729 + if comp is None:
730 + await finish_discover_job(j["id"], ok=False, error="company not found", attempts=j["attempts"], max_attempts=j["max_attempts"])
731 + continue
732 + return comp, j
703 733 async with transaction() as conn:
704 − if company_slug:
734 + row = await fetch_one(conn, """
735 + update companies set onboarding_status = 'discovering', updated_at = now()
736 + where id = (select id from companies where onboarding_status = 'pending' order by importance desc, created_at limit 1 for update skip locked)
737 + returning *""")
738 + if row is None:
739 + async with lock:
740 + budget["left"] += 1
741 + return None
742 + return row, None
743 +
744 + async def worker_loop() -> None:
745 + while True:
746 + target = await claim_one()
747 + if target is None:
748 + return
749 + async with lock:
750 + stats["claimed"] += 1
751 + try:
752 + await process(*target)
753 + except Exception: # one company must never stop the pool
754 + log.exception("onboarding worker failed", extra={"company": target[0].get("slug")})
755 +
756 + try:
757 + if company_slug:
758 + async with transaction() as conn:
705 759 comp = await fetch_one(conn, "select * from companies where slug = :s or id = :s or canonical_domain = :s", s=company_slug)
706 − if comp is None:
707 − raise LookupError(f"company {company_slug!r} not found")
708 − targets.append((comp, None))
709 − if not company_slug:
710 − jobs = await claim_discover_jobs(limit, worker)
711 − for j in jobs:
712 − cid = (j.get("payload") or {}).get("company_id") or j["key"].removeprefix("discover:")
713 − async with transaction() as conn:
714 − comp = await fetch_one(conn, "select * from companies where id = :id or slug = :id", id=cid)
715 − if comp is None:
716 − await finish_discover_job(j["id"], ok=False, error="company not found", attempts=j["attempts"], max_attempts=j["max_attempts"])
717 − continue
718 − targets.append((comp, j))
719 − if len(targets) < limit:
720 − async with transaction() as conn:
721 − rows = await fetch_all(conn, """
722 − update companies set onboarding_status = 'discovering', updated_at = now()
723 − where id in (select id from companies where onboarding_status = 'pending' and id <> all(cast(:skip as text[]))
724 − order by importance desc, created_at limit :n for update skip locked)
725 − returning *""", n=limit - len(targets), skip=[c["id"] for c, _ in targets] or ["-"])
726 − targets.extend((r, None) for r in rows)
727 − stats["claimed"] = len(targets)
728 − sem = asyncio.Semaphore(max(1, concurrency))
729 −
730 − async def one(comp: dict[str, Any], job: dict[str, Any] | None) -> None:
731 − async with sem:
732 − res = await discover_company(comp, fetcher=fetcher, dry_run=dry_run, fetch_now=fetch_now)
733 − stats[str(res.status)] = stats.get(str(res.status), 0) + 1
734 − stats["sensors"] += len(res.sensors)
735 − if job is not None:
736 − await finish_discover_job(job["id"], ok=res.status == OnboardingStatus.ACTIVE, error=res.error, attempts=job["attempts"], max_attempts=job["max_attempts"])
737 − log.info("company discovered", extra={"company": comp.get("slug"), "status": str(res.status), "sensors": len(res.sensors), "requests": res.requests,
738 − "ms": res.duration_ms, "error": res.error})
739 − if fetch_now and not dry_run and res.sensors:
740 − from companyatlas.services.pipeline import run_sensor_ids
741 −
742 − await run_sensor_ids([s["id"] for s in res.sensors], fetcher=fetcher, worker=worker)
743 −
744 − await asyncio.gather(*(one(c, j) for c, j in targets))
760 + if comp is None:
761 + raise LookupError(f"company {company_slug!r} not found")
762 + stats["claimed"] = 1
763 + await process(comp, None)
764 + else:
765 + await asyncio.gather(*(worker_loop() for _ in range(max(1, concurrency))))
745 766 finally:
746 767 if own_fetcher:
747 768 await fetcher.close()
748 769