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

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:
NixReaper
2026-09-05 23:08:37 -05:00
co-authored by Claude Sonnet 5
parent a52ac676ff
commit b05d34db5f
2 changed files with 133 additions and 13 deletions
+29 -13
View File
@@ -87,25 +87,41 @@ 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
if enqueued % 5000 == 0:
logger.info(
"bulk_import_progress",
entity_type=entity_type,
enqueued=enqueued,
total=len(ids),
)
async with lock:
enqueued += 1
if enqueued % 5000 == 0:
logger.info(
"bulk_import_progress",
entity_type=entity_type,
enqueued=enqueued,
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)
+104
View File
@@ -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