feat: batch processing

Author: m0wer

Diff for: backend/app/routers/feed.py

@@ -1,4 +1,6 @@
+import asyncio
 from pydantic.networks import HttpUrl
+from datetime import datetime, timedelta, timezone
 import re
 from fastapi import Request
 import json
@@ -8,8 +10,10 @@
 from app.models.feed import Feed, generate_feed, parse_feed, UpstreamError
 from app.models.user import User
 from app.recommend import filter_articles
+from app.tasks import fetch_feed_batch
 from .common import get_engine
 from fastapi import HTTPException
+from fastapi import BackgroundTasks
 
 # from fastapi_cache.coder import PickleCoder
 # from fastapi_cache.decorator import cache
@@ -21,35 +25,33 @@
 )
 
 
+FEED_REFRESH_INTERVAL = timedelta(days=1)  # Adjust as needed
+
+
 @router.get("/{user_id}/{feed_url:path}")
-# @cache(expire=300, coder=PickleCoder)
 async def get_feed(
-    request: Request, user_id: str, feed_url: HttpUrl, engine=Depends(get_engine)
+    request: Request,
+    user_id: str,
+    feed_url: HttpUrl,
+    background_tasks: BackgroundTasks,
+    engine=Depends(get_engine),
 ) -> str:
-    """Get filtered feed."""
     if feed_url.host and re.match(
         r"^(25[0-5]|(2[0-4]|1\d|[1-9]|)\d)\.?\b", feed_url.host
     ):
         raise HTTPException(status_code=422, detail="Invalid URL")
-    if request.query_params:
-        feed_url = f"{feed_url}?"  # type: ignore
-        for key, value in request.query_params.items():
-            feed_url = f"{feed_url}&{key}={value}"  # type: ignore
+
     with Session(engine, autoflush=False) as session:
+        # User handling (same as before)
         try:
             user: User = session.exec(select(User).where(User.id == user_id)).one()
         except NoResultFound:
             logger.info(f"User {user_id} not found in database, creating new user")
             user = User(id=user_id)
             session.add(user)
-            try:
-                session.commit()
-            except Exception as e:
-                # might happen if the user was created before by another thread
-                logger.warning(f"Failed to add user {user_id} to database: {e}")
-                session.rollback()
-                user = session.exec(select(User).where(User.id == user_id)).one()
+            session.commit()
 
+        # Feed handling
         try:
             feed: Feed = session.exec(
                 select(Feed).where(Feed.url == str(feed_url))
@@ -70,9 +72,23 @@ async def get_feed(
                 logger.warning(f"Failed to add feed {feed_url} to database: {e}")
                 session.rollback()
                 feed = session.exec(select(Feed).where(Feed.url == str(feed_url))).one()
+            session.add(feed)
+            session.commit()
+
         if feed not in user.feeds:
             user.feeds.append(feed)
-        session.commit()
+            session.commit()
+
+        # Check if feed needs refreshing
+        now = datetime.now(timezone.utc)
+        if (
+            feed.updated_at is None
+            or (now - feed.updated_at.replace(tzinfo=timezone.utc))
+            > FEED_REFRESH_INTERVAL
+        ):
+            logger.info(f"Feed {feed_url} needs refreshing")
+            await asyncio.run(fetch_feed_batch([feed.id]))
+            session.refresh(feed)
 
         articles = feed.articles[-30:]
 

Diff for: backend/app/tasks.py

@@ -4,8 +4,8 @@
 from sqlmodel import create_engine, Session, select
 from datetime import datetime, timezone, timedelta
 from loguru import logger
-from redis import Redis
-from rq import Queue
+from redis import Redis  # type: ignore
+from rq import Queue, Retry
 
 from app.models.article import Article
 from app.models.feed import Feed, parse_feed, generate_feed
@@ -25,32 +25,6 @@
 gpu_queue = Queue("gpu", connection=redis_conn)
 
 
-def fetch_feed(feed_id):
-    with Session(ENGINE) as session:
-        feed = session.get(Feed, feed_id)
-        if not feed:
-            logger.error(f"Feed {feed_id} not found")
-            return
-
-        try:
-            parsed_feed = asyncio.run(parse_feed(feed.url))
-            for article in parsed_feed.articles:
-                existing_article = session.exec(
-                    select(Article).where(Article.url == article.url)
-                ).first()
-                if not existing_article:
-                    article.feed = feed
-                    session.add(article)
-                    session.commit()
-                    gpu_queue.enqueue(compute_article_embedding, article.id)
-
-            logger.info(
-                f"Fetched {len(parsed_feed.articles)} articles for feed {feed_id}"
-            )
-        except Exception as e:
-            logger.error(f"Error fetching feed {feed_id}: {e}")
-
-
 def compute_article_embedding(article_id):
     with Session(ENGINE) as session:
         article = session.get(Article, article_id)
@@ -87,7 +61,7 @@ def remove_old_embeddings():
         old_articles = session.exec(
             select(Article)
             .where(Article.updated < one_month_ago)
-            .where(Article.embedding is not None)
+            .where(Article.embedding != None)  # noqa: E711
         ).all()
 
         for article in old_articles:
@@ -97,12 +71,86 @@ def remove_old_embeddings():
         logger.info(f"Removed embeddings from {len(old_articles)} old articles")
 
 
+BATCH_SIZE = int(os.getenv("FEED_FETCH_BATCH_SIZE", "50"))
+
+
+async def fetch_feed_batch(feed_ids):
+    async def fetch_single_feed(feed):
+        try:
+            return await parse_feed(feed.url)
+        except Exception as e:
+            logger.error(f"Error fetching feed {feed.id}: {e}")
+            return None
+
+    with Session(ENGINE) as session:
+        feeds = session.exec(select(Feed).where(Feed.id.in_(feed_ids))).all()
+        tasks = [asyncio.create_task(fetch_single_feed(feed)) for feed in feeds]
+        results = await asyncio.gather(*tasks)
+
+        new_articles = []
+        for feed, parsed_feed in zip(feeds, results):
+            if parsed_feed is None:
+                continue
+
+            for article in parsed_feed.articles:
+                existing_article = session.exec(
+                    select(Article).where(Article.url == article.url)
+                ).first()
+                if not existing_article:
+                    article.feed = feed
+                    session.add(article)
+                    new_articles.append(article)
+
+            feed.updated_at = datetime.now(timezone.utc)
+
+        session.commit()
+
+        if new_articles:
+            new_article_ids = [article.id for article in new_articles]
+            enqueue_gpu_task(compute_embeddings_batch, new_article_ids)
+
+        logger.info(
+            f"Fetched {len(feeds)} feeds, added {len(new_articles)} new articles"
+        )
+
+
+def compute_embeddings_batch(article_ids):
+    with Session(ENGINE) as session:
+        articles = session.exec(
+            select(Article).where(Article.id.in_(article_ids))
+        ).all()
+        articles_to_embed = [
+            article for article in articles if article.embedding is None
+        ]
+
+        if not articles_to_embed:
+            return
+
+        try:
+            compute_embeddings(articles_to_embed)
+            session.commit()
+            logger.info(f"Computed embeddings for {len(articles_to_embed)} articles")
+        except Exception as e:
+            logger.error(f"Error computing embeddings for articles: {e}")
+
+
 def fetch_all_feeds():
     with Session(ENGINE) as session:
-        feeds = session.exec(select(Feed)).all()
-        for feed in feeds:
-            low_queue.enqueue(fetch_feed, feed.id)
-    logger.info(f"Enqueued fetch tasks for {len(feeds)} feeds")
+        one_month_ago = datetime.now(timezone.utc) - timedelta(days=30)
+        active_feeds = session.exec(
+            select(Feed)
+            .join(User.feeds)
+            .where(User.last_request > one_month_ago)
+            .distinct()
+        ).all()
+
+        for i in range(0, len(active_feeds), BATCH_SIZE):
+            batch = active_feeds[i : i + BATCH_SIZE]
+            asyncio.run(fetch_feed_batch([feed.id for feed in batch]))
+
+    logger.info(
+        f"Processed {len(active_feeds)} active feeds in batches of {BATCH_SIZE}"
+    )
 
 
 def log_user_action(user_id: str, article_id: int, link_url: str):
@@ -123,7 +171,7 @@ def log_user_action(user_id: str, article_id: int, link_url: str):
         if article not in user.articles:
             user.articles.append(article)
             session.add(user)
-            medium_queue.enqueue(recompute_user_clusters, user.id)
+            enqueue_medium_priority(recompute_user_clusters, user.id)
 
         session.commit()
     logger.info(f"Logged action for user {user_id}, article {article_id}")
@@ -151,12 +199,24 @@ def generate_filtered_feed(feed_id: int, user_id: str):
 
 # Helper functions to enqueue tasks
 def enqueue_low_priority(func, *args, **kwargs):
-    return low_queue.enqueue(func, *args, **kwargs)
+    return low_queue.enqueue(
+        func, args=args, kwargs=kwargs, retry=Retry(max=3), timeout=300
+    )
 
 
 def enqueue_medium_priority(func, *args, **kwargs):
-    return medium_queue.enqueue(func, *args, **kwargs)
+    return medium_queue.enqueue(
+        func, args=args, kwargs=kwargs, retry=Retry(max=3), timeout=300
+    )
 
 
 def enqueue_high_priority(func, *args, **kwargs):
-    return high_queue.enqueue(func, *args, **kwargs)
+    return high_queue.enqueue(
+        func, args=args, kwargs=kwargs, retry=Retry(max=2), timeout=10
+    )
+
+
+def enqueue_gpu_task(func, *args, **kwargs):
+    return gpu_queue.enqueue(
+        func, args=args, kwargs=kwargs, retry=Retry(max=3), timeout=600
+    )