Fix bulk_import.py stalling on real-scale exports (1M+ ids)
CI / lint (push) Canceled after 0s
CI / typecheck (push) Canceled after 0s
CI / unit-tests (push) Canceled after 0s
CI / integration-tests (push) Canceled after 0s
CI / container-build (push) Canceled after 0s
CI / dependency-audit (push) Canceled after 0s
CI / lint (push) Canceled after 0s
CI / typecheck (push) Canceled after 0s
CI / unit-tests (push) Canceled after 0s
CI / integration-tests (push) Canceled after 0s
CI / container-build (push) Canceled after 0s
CI / dependency-audit (push) Canceled after 0s
asyncio.gather(*(_worker(id) for id in ids)) eagerly wraps every argument into a Task before awaiting any of them -- with a semaphore bounding concurrent *execution* but not concurrent Task *creation*, running this against the real movie export (1,240,787 ids) meant constructing 1.24 million live Task + coroutine objects up front, which looked indistinguishable from hanging for minutes before any work could start, confirmed live in production tonight. Replaced with a fixed pool of _CONCURRENCY (50) consumer tasks pulling from a plain asyncio.Queue pre-loaded with every id -- only 50 Task objects ever exist regardless of catalog size. Verified: 500K enqueues against no-op stubs now complete in 0.59s (previously: the setup loop alone dominated), and the real bottleneck is correctly back to the per-id DB round-trip + Redis push, gated by the pool of 50 -- not Python-level Task scheduling overhead. Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
co-authored by
Claude Sonnet 5
parent
a52ac676ff
commit
b05d34db5f
@@ -87,13 +87,29 @@ async def _bulk_enqueue(entity_type: str, *, dry_run: bool, include_seasons: boo
|
||||
return
|
||||
|
||||
queue = await get_arq_pool()
|
||||
semaphore = asyncio.Semaphore(_CONCURRENCY)
|
||||
# A fixed pool of _CONCURRENCY consumer tasks pulling from a plain
|
||||
# asyncio.Queue -- NOT asyncio.gather() over one Task per id. gather()
|
||||
# eagerly wraps every argument into a Task before awaiting any of them,
|
||||
# so with a catalog in the millions that means millions of live Task +
|
||||
# coroutine objects created up front regardless of how tightly a
|
||||
# semaphore bounds actual concurrent execution -- a real bottleneck
|
||||
# (and memory risk) discovered running this against the real 1.24M-
|
||||
# movie export, where it looked hung for minutes just getting through
|
||||
# that setup before any work could start.
|
||||
id_queue: asyncio.Queue[str] = asyncio.Queue()
|
||||
for external_id in ids:
|
||||
id_queue.put_nowait(external_id)
|
||||
|
||||
enqueued = 0
|
||||
lock = asyncio.Lock()
|
||||
|
||||
async def _worker(external_id: str) -> None:
|
||||
async def _consumer() -> None:
|
||||
nonlocal enqueued
|
||||
async with semaphore:
|
||||
while True:
|
||||
try:
|
||||
external_id = id_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
return
|
||||
await _enqueue_one(queue, entity_type, external_id, include_seasons=include_seasons)
|
||||
async with lock:
|
||||
enqueued += 1
|
||||
@@ -105,7 +121,7 @@ async def _bulk_enqueue(entity_type: str, *, dry_run: bool, include_seasons: boo
|
||||
total=len(ids),
|
||||
)
|
||||
|
||||
await asyncio.gather(*(_worker(external_id) for external_id in ids))
|
||||
await asyncio.gather(*(_consumer() for _ in range(_CONCURRENCY)))
|
||||
logger.info("bulk_import_complete", entity_type=entity_type, enqueued=enqueued)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
"""Unit tests for app/scripts/bulk_import.py's enqueue logic.
|
||||
|
||||
Exercises _bulk_enqueue's bounded worker-pool directly (network, arq, and
|
||||
the DB are all monkeypatched out) -- the property under test is that every
|
||||
id from the export gets enqueued exactly once, with no duplicates and none
|
||||
skipped, which is the correctness risk of moving off asyncio.gather(*ids)
|
||||
onto a fixed pool of consumers pulling from an asyncio.Queue (done after
|
||||
the gather-based version was found, live, to effectively stall for
|
||||
minutes on Task-creation overhead alone against a real 1.24M-id export)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from app.scripts import bulk_import
|
||||
|
||||
|
||||
class _FakeQueue:
|
||||
def __init__(self) -> None:
|
||||
self.enqueued: list[tuple] = []
|
||||
|
||||
async def enqueue_job(self, task_name: str, *args):
|
||||
self.enqueued.append((task_name, args))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_enqueue_enqueues_every_id_exactly_once(monkeypatch):
|
||||
ids = [str(n) for n in range(1, 501)]
|
||||
fake_queue = _FakeQueue()
|
||||
|
||||
async def fake_fetch(entity_type, *, client):
|
||||
return ids
|
||||
|
||||
async def fake_get_or_create_ingestion_job(job_type, entity_type, external_id):
|
||||
return uuid.uuid4()
|
||||
|
||||
async def fake_get_arq_pool():
|
||||
return fake_queue
|
||||
|
||||
monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch)
|
||||
monkeypatch.setattr(
|
||||
bulk_import, "get_or_create_ingestion_job", fake_get_or_create_ingestion_job
|
||||
)
|
||||
monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool)
|
||||
|
||||
await bulk_import._bulk_enqueue("movie", dry_run=False, include_seasons=False)
|
||||
|
||||
enqueued_ids = sorted(int(args[1]) for _, args in fake_queue.enqueued)
|
||||
assert enqueued_ids == list(range(1, 501))
|
||||
assert len(fake_queue.enqueued) == 500
|
||||
assert all(task_name == "ingest_movie_task" for task_name, _ in fake_queue.enqueued)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_enqueue_dry_run_enqueues_nothing(monkeypatch):
|
||||
ids = [str(n) for n in range(1, 51)]
|
||||
fake_queue = _FakeQueue()
|
||||
called = False
|
||||
|
||||
async def fake_fetch(entity_type, *, client):
|
||||
return ids
|
||||
|
||||
async def fake_get_arq_pool():
|
||||
nonlocal called
|
||||
called = True
|
||||
return fake_queue
|
||||
|
||||
monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch)
|
||||
monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool)
|
||||
|
||||
await bulk_import._bulk_enqueue("collection", dry_run=True, include_seasons=False)
|
||||
|
||||
assert called is False
|
||||
assert fake_queue.enqueued == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_enqueue_passes_include_seasons_flag_for_series(monkeypatch):
|
||||
fake_queue = _FakeQueue()
|
||||
|
||||
async def fake_fetch(entity_type, *, client):
|
||||
return ["1396"]
|
||||
|
||||
async def fake_get_or_create_ingestion_job(job_type, entity_type, external_id):
|
||||
return uuid.uuid4()
|
||||
|
||||
async def fake_get_arq_pool():
|
||||
return fake_queue
|
||||
|
||||
monkeypatch.setattr(bulk_import, "fetch_daily_export_ids", fake_fetch)
|
||||
monkeypatch.setattr(
|
||||
bulk_import, "get_or_create_ingestion_job", fake_get_or_create_ingestion_job
|
||||
)
|
||||
monkeypatch.setattr(bulk_import, "get_arq_pool", fake_get_arq_pool)
|
||||
|
||||
await bulk_import._bulk_enqueue("series", dry_run=False, include_seasons=True)
|
||||
|
||||
assert len(fake_queue.enqueued) == 1
|
||||
task_name, args = fake_queue.enqueued[0]
|
||||
assert task_name == "ingest_series_task"
|
||||
assert args[1] == "1396"
|
||||
assert args[2] is True
|
||||
Reference in New Issue
Block a user