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 1 mo ago (Sep 12, 2026) parent f09c483 Browse files

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	def register(app: typer.Typer) -> None:	27	def register(app: typer.Typer) -> None:
28	@app.command()	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	concurrency: int \| None = None, fetch_now: bool = False, dry_run: bool = False) -> None:	30	concurrency: int \| None = None, fetch_now: bool = False, dry_run: bool = False) -> None:
31	"""Discover surfaces and create sensors for pending companies (or one company)."""	31	"""Discover surfaces and create sensors for pending companies (or one company)."""
32	from companyatlas.services.discovery import onboard_pending	32	from companyatlas.services.discovery import onboard_pending

modified src/companyatlas/config.py +1 −0

				`@@ -50,6 +50,7 @@ class Settings(BaseSettings):`
50	discovery_min_confidence: float = Field(0.55, alias="CA_DISCOVERY_MIN_CONFIDENCE")	50	discovery_min_confidence: float = Field(0.55, alias="CA_DISCOVERY_MIN_CONFIDENCE")
51	discovery_max_sensors_per_company: int = Field(40, alias="CA_DISCOVERY_MAX_SENSORS")	51	discovery_max_sensors_per_company: int = Field(40, alias="CA_DISCOVERY_MAX_SENSORS")
52	onboarding_concurrency: int = Field(12, alias="CA_ONBOARDING_CONCURRENCY")	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	# ---------------------------------------------------------------- scheduling (seconds)	55	# ---------------------------------------------------------------- scheduling (seconds)
55	scheduler_tick_s: int = Field(15, alias="CA_SCHEDULER_TICK_S")	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	accept: str \| None = None) -> FetchResult \| None:	152	accept: str \| None = None) -> FetchResult \| None:
153	res.requests += 1	153	res.requests += 1
154	try:	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	except NotModified:	156	except NotModified:
157	return None	157	return None
158	except (FetchError, BlockedError) as exc:	158	except (FetchError, BlockedError) as exc:
				`@@ -259,7 +259,7 @@ async def _phase_homepage(fetcher: Fetcher, company: dict[str, Any], res: Discov`
259	break	259	break
260	res.requests += 1	260	res.requests += 1
261	try:	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	except BlockedError as exc:	263	except BlockedError as exc:
264	res.notes.append(f"homepage blocked: {exc.failure}")	264	res.notes.append(f"homepage blocked: {exc.failure}")
265	res.error = f"{exc.failure}: {exc}"[:300]	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	run_at = now() + make_interval(mins => :mins) where id = :id""", id=job_id, e=(error or "")[:500], mins=30 * attempts)	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	fetch_now: bool = False, dry_run: bool = False, worker: str = "onboard") -> dict[str, Any]:	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	concurrency = concurrency or settings.onboarding_concurrency	697	concurrency = concurrency or settings.onboarding_concurrency
696	own_fetcher = fetcher is None	698	own_fetcher = fetcher is None
697	fetcher = fetcher or Fetcher()	699	fetcher = fetcher or Fetcher()
698	if own_fetcher:	700	if own_fetcher:
699	await fetcher.open()	701	await fetcher.open()
700	stats = {"claimed": 0, "active": 0, "failed": 0, "no_website": 0, "sensors": 0}	702	stats = {"claimed": 0, "active": 0, "failed": 0, "no_website": 0, "sensors": 0}
701	− try:	703	+ lock = asyncio.Lock()
702	− targets: list[tuple[dict[str, Any], dict[str, Any] \| None]] = []	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	async with transaction() as conn:	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	comp = await fetch_one(conn, "select * from companies where slug = :s or id = :s or canonical_domain = :s", s=company_slug)	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:	760	+ if comp is None:
707	− raise LookupError(f"company {company_slug!r} not found")	761	+ raise LookupError(f"company {company_slug!r} not found")
708	− targets.append((comp, None))	762	+ stats["claimed"] = 1
709	− if not company_slug:	763	+ await process(comp, None)
710	− jobs = await claim_discover_jobs(limit, worker)	764	+ else:
711	− for j in jobs:	765	+ await asyncio.gather(*(worker_loop() for _ in range(max(1, concurrency))))
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))
745	finally:	766	finally:
746	if own_fetcher:	767	if own_fetcher:
747	await fetcher.close()	768	await fetcher.close()
748		769

spb/company-atlas

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

Theme