fix(ml): per-task async engine for recompute_centroid (#881)
recompute_centroid + recompute_centroids were the only tasks still using the process-wide singleton extensions.get_session() under asyncio.run(). The async engine's asyncpg pool is bound to the loop it was created on; each Celery task runs a fresh asyncio.run() loop, so after the first invocation the cached engine handed loop-A connections to loop B and raised "Future attached to a different loop" — every recompute after the first in a worker process failed (~35ms, fails on first DB await). Convert both to the established per-task async_session_factory() pattern (NullPool engine created + disposed inside the task's own loop), matching scan/download/admin tasks. No get_session usages remain in tasks/. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+21
-9
@@ -491,15 +491,22 @@ def _confidence_for_tag(session, tag, preds: dict) -> float | None:
|
|||||||
def recompute_centroid(self, tag_id: int) -> bool:
|
def recompute_centroid(self, tag_id: int) -> bool:
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
from ..extensions import get_session
|
|
||||||
from ..services.ml.centroids import CentroidService
|
from ..services.ml.centroids import CentroidService
|
||||||
|
from ._async_session import async_session_factory
|
||||||
|
|
||||||
async def _run() -> bool:
|
async def _run() -> bool:
|
||||||
async with get_session() as session:
|
# Per-task NullPool engine bound to THIS asyncio.run loop — the shared
|
||||||
svc = CentroidService(session)
|
# process-wide engine reuses connections across loops and raises
|
||||||
result = await svc.recompute_for_tag(tag_id)
|
# "Future attached to a different loop" on every call after the first.
|
||||||
await session.commit()
|
async_factory, async_engine = async_session_factory()
|
||||||
return result
|
try:
|
||||||
|
async with async_factory() as session:
|
||||||
|
svc = CentroidService(session)
|
||||||
|
result = await svc.recompute_for_tag(tag_id)
|
||||||
|
await session.commit()
|
||||||
|
return result
|
||||||
|
finally:
|
||||||
|
await async_engine.dispose()
|
||||||
|
|
||||||
return asyncio.run(_run())
|
return asyncio.run(_run())
|
||||||
|
|
||||||
@@ -515,12 +522,17 @@ def recompute_centroids(self) -> int:
|
|||||||
"""Daily: find drifted centroids, enqueue recompute_centroid for each."""
|
"""Daily: find drifted centroids, enqueue recompute_centroid for each."""
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
from ..extensions import get_session
|
|
||||||
from ..services.ml.centroids import CentroidService
|
from ..services.ml.centroids import CentroidService
|
||||||
|
from ._async_session import async_session_factory
|
||||||
|
|
||||||
async def _list() -> list[int]:
|
async def _list() -> list[int]:
|
||||||
async with get_session() as session:
|
# Per-task NullPool engine bound to this loop (see recompute_centroid).
|
||||||
return await CentroidService(session).list_drifted()
|
async_factory, async_engine = async_session_factory()
|
||||||
|
try:
|
||||||
|
async with async_factory() as session:
|
||||||
|
return await CentroidService(session).list_drifted()
|
||||||
|
finally:
|
||||||
|
await async_engine.dispose()
|
||||||
|
|
||||||
drifted = asyncio.run(_list())
|
drifted = asyncio.run(_list())
|
||||||
for tid in drifted:
|
for tid in drifted:
|
||||||
|
|||||||
Reference in New Issue
Block a user