xueyiming vor 4 Tagen
Ursprung
Commit
278b988b5d

+ 15 - 0
supply_infra/db/repositories/video_discovery_repo.py

@@ -349,6 +349,21 @@ class VideoDiscoveryRepository(BaseRepository[VideoDiscoveryRun]):
             stmt = stmt.limit(limit)
             stmt = stmt.limit(limit)
         return list(self.session.scalars(stmt).all())
         return list(self.session.scalars(stmt).all())
 
 
+    def count_passed_videos(self, biz_dt: str) -> int:
+        """统计业务日内 primary / backup 的唯一视频数。"""
+        stmt = (
+            select(func.count(func.distinct(VideoDiscoveryCandidate.aweme_id)))
+            .join(
+                VideoDiscoveryRun,
+                VideoDiscoveryRun.run_id == VideoDiscoveryCandidate.run_id,
+            )
+            .where(VideoDiscoveryRun.biz_dt == biz_dt)
+            .where(VideoDiscoveryCandidate.decision_bucket.in_(_PUBLISHABLE_BUCKETS))
+            .where(VideoDiscoveryCandidate.aweme_id.is_not(None))
+            .where(func.trim(VideoDiscoveryCandidate.aweme_id) != "")
+        )
+        return int(self.session.scalar(stmt) or 0)
+
     def mark_candidates_aigc_plans(
     def mark_candidates_aigc_plans(
         self,
         self,
         candidate_ids: list[int],
         candidate_ids: list[int],

+ 71 - 36
supply_infra/scheduler/jobs/discover_videos_from_demands.py

@@ -21,11 +21,13 @@ from agents.find_agent.demand_run import (
 )
 )
 from supply_infra.config import get_infra_settings
 from supply_infra.config import get_infra_settings
 from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
 from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
+from supply_infra.db.repositories.video_discovery_repo import VideoDiscoveryRepository
 from supply_infra.db.session import get_session
 from supply_infra.db.session import get_session
 
 
 logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
 
 
 _DEFAULT_WORKERS = 1
 _DEFAULT_WORKERS = 1
+_DAILY_PASSED_VIDEO_LIMIT = 200
 
 
 
 
 def _resolve_biz_dt(biz_dt: str | None) -> str:
 def _resolve_biz_dt(biz_dt: str | None) -> str:
@@ -96,6 +98,11 @@ def process_single_discover(
         }
         }
 
 
 
 
+def _count_passed_videos(biz_dt: str) -> int:
+    with get_session() as session:
+        return VideoDiscoveryRepository(session).count_passed_videos(biz_dt)
+
+
 def discover_videos_from_demands(
 def discover_videos_from_demands(
     biz_dt: str | None = None,
     biz_dt: str | None = None,
     *,
     *,
@@ -144,8 +151,10 @@ def discover_videos_from_demands(
     if limit is not None and limit >= 0:
     if limit is not None and limit >= 0:
         contexts = contexts[: int(limit)]
         contexts = contexts[: int(limit)]
 
 
+    passed_videos = _count_passed_videos(resolved_biz_dt)
+
     logger.info(
     logger.info(
-        "discover_videos_from_demands start: biz_dt=%s batch_run_id=%s workers=%s top_limit=%s offset=%s pending=%s skipped=%s",
+        "discover_videos_from_demands start: biz_dt=%s batch_run_id=%s workers=%s top_limit=%s offset=%s pending=%s skipped=%s passed_videos=%s passed_video_limit=%s",
         resolved_biz_dt,
         resolved_biz_dt,
         batch_run_id,
         batch_run_id,
         workers,
         workers,
@@ -153,6 +162,8 @@ def discover_videos_from_demands(
         offset,
         offset,
         len(contexts),
         len(contexts),
         preload_stats.get("skipped_already_done", 0),
         preload_stats.get("skipped_already_done", 0),
+        passed_videos,
+        _DAILY_PASSED_VIDEO_LIMIT,
     )
     )
 
 
     result: dict[str, Any] = {
     result: dict[str, Any] = {
@@ -170,9 +181,13 @@ def discover_videos_from_demands(
         "skipped": 0,
         "skipped": 0,
         "failed": 0,
         "failed": 0,
         "errors": [],
         "errors": [],
+        "passed_videos": passed_videos,
+        "passed_video_limit": _DAILY_PASSED_VIDEO_LIMIT,
+        "stopped_by_passed_video_limit": passed_videos
+        >= _DAILY_PASSED_VIDEO_LIMIT,
     }
     }
 
 
-    if not contexts:
+    if not contexts or passed_videos >= _DAILY_PASSED_VIDEO_LIMIT:
         finished_at = datetime.now()
         finished_at = datetime.now()
         result["finished_at"] = finished_at.isoformat()
         result["finished_at"] = finished_at.isoformat()
         result["duration_seconds"] = round((finished_at - started_at).total_seconds(), 2)
         result["duration_seconds"] = round((finished_at - started_at).total_seconds(), 2)
@@ -183,41 +198,61 @@ def discover_videos_from_demands(
     result["workers"] = worker_count
     result["workers"] = worker_count
 
 
     with ThreadPoolExecutor(max_workers=worker_count) as executor:
     with ThreadPoolExecutor(max_workers=worker_count) as executor:
-        futures = [
-            executor.submit(process_single_discover, ctx, force=force)
-            for ctx in contexts
-        ]
-        for future in as_completed(futures):
-            try:
-                item_result = future.result()
-            except Exception as exc:
-                logger.exception(
-                    "discover videos worker 出现未捕获错误: biz_dt=%s",
-                    resolved_biz_dt,
-                )
-                result["failed"] += 1
-                result["processed"] += 1
-                result["errors"].append({"error": str(exc)})
-                continue
-
-            result["processed"] += 1
-            if item_result.get("skipped"):
-                result["skipped"] += 1
-                continue
-            if item_result.get("success"):
-                result["succeeded"] += 1
-                continue
-
-            result["failed"] += 1
-            result["errors"].append(
-                {
-                    "demand_grade_id": item_result.get("demand_grade_id"),
-                    "demand_name": item_result.get("demand_name"),
-                    "video_count": item_result.get("video_count"),
-                    "run_id": item_result.get("run_id"),
-                    "error": item_result.get("error"),
-                }
+        next_context = 0
+        while next_context < len(contexts):
+            remaining_slots = _DAILY_PASSED_VIDEO_LIMIT - passed_videos
+            if remaining_slots <= 0:
+                result["stopped_by_passed_video_limit"] = True
+                break
+
+            batch_size = min(
+                worker_count,
+                remaining_slots,
+                len(contexts) - next_context,
             )
             )
+            batch = contexts[next_context : next_context + batch_size]
+            next_context += batch_size
+            futures = [
+                executor.submit(process_single_discover, ctx, force=force)
+                for ctx in batch
+            ]
+            for future in as_completed(futures):
+                try:
+                    item_result = future.result()
+                except Exception as exc:
+                    logger.exception(
+                        "discover videos worker 出现未捕获错误: biz_dt=%s",
+                        resolved_biz_dt,
+                    )
+                    result["failed"] += 1
+                    result["processed"] += 1
+                    result["errors"].append({"error": str(exc)})
+                    continue
+
+                result["processed"] += 1
+                if item_result.get("skipped"):
+                    result["skipped"] += 1
+                    continue
+                if item_result.get("success"):
+                    result["succeeded"] += 1
+                    continue
+
+                result["failed"] += 1
+                result["errors"].append(
+                    {
+                        "demand_grade_id": item_result.get("demand_grade_id"),
+                        "demand_name": item_result.get("demand_name"),
+                        "video_count": item_result.get("video_count"),
+                        "run_id": item_result.get("run_id"),
+                        "error": item_result.get("error"),
+                    }
+                )
+
+            passed_videos = _count_passed_videos(resolved_biz_dt)
+            result["passed_videos"] = passed_videos
+            if passed_videos >= _DAILY_PASSED_VIDEO_LIMIT:
+                result["stopped_by_passed_video_limit"] = True
+                break
 
 
     finished_at = datetime.now()
     finished_at = datetime.now()
     result["finished_at"] = finished_at.isoformat()
     result["finished_at"] = finished_at.isoformat()

+ 40 - 0
tests/supply_infra/scheduler/test_discover_videos_from_demands.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from supply_infra.scheduler.jobs.discover_videos_from_demands import (
+    discover_videos_from_demands,
+)
+
+
+@patch(
+    "supply_infra.scheduler.jobs.discover_videos_from_demands.process_single_discover"
+)
+@patch("supply_infra.scheduler.jobs.discover_videos_from_demands._count_passed_videos")
+@patch(
+    "supply_infra.scheduler.jobs.discover_videos_from_demands.filter_pending_contexts"
+)
+@patch(
+    "supply_infra.scheduler.jobs.discover_videos_from_demands.list_find_demand_contexts"
+)
+def test_stops_discovery_after_200_passed_videos(
+    mock_list_contexts,
+    mock_filter_contexts,
+    mock_count_passed,
+    mock_process,
+) -> None:
+    contexts = [object(), object()]
+    mock_list_contexts.return_value = ("20260727", contexts)
+    mock_filter_contexts.return_value = (
+        contexts,
+        {"total_loaded": 2, "skipped_already_done": 0},
+    )
+    mock_count_passed.side_effect = [199, 200]
+    mock_process.return_value = {"success": True, "skipped": False}
+
+    result = discover_videos_from_demands("20260727", workers=2)
+
+    assert mock_process.call_count == 1
+    assert result["processed"] == 1
+    assert result["passed_videos"] == 200
+    assert result["stopped_by_passed_video_limit"] is True