|
|
@@ -2,17 +2,22 @@
|
|
|
定时任务:对 multi_demand_pool_di 需求池中的需求循环打分级。
|
|
|
|
|
|
循环控制在本文件(调度任务侧),而不是让 agent 在一次运行内自行分页遍历——
|
|
|
-每一批只挑选少量(默认 20 个)尚未分级的需求词,调用一次 demand_grade_agent
|
|
|
-并等待其运行结束(agent 内部会调用 batch_save_demand_grades 落库),再重新
|
|
|
-查询"已分级"集合、找出下一批新词,如此循环直到没有更多待分级需求或达到本次
|
|
|
-运行的批次上限,从而避免单次 agent 会话上下文无限增长。
|
|
|
+每一轮并行启动多个 worker(默认 5 个线程),每个 worker 处理一批互不重叠的
|
|
|
+需求词(默认每批 20 个),调用一次 demand_grade_agent 并等待其运行结束
|
|
|
+(agent 内部会调用 batch_save_demand_grades 落库),再重新查询"已分级"集合、
|
|
|
+分配下一批新词,如此循环直到没有更多待分级需求或达到本次运行的批次上限,
|
|
|
+从而避免单次 agent 会话上下文无限增长,同时提升吞吐。
|
|
|
"""
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import logging
|
|
|
+import math
|
|
|
+from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
from datetime import datetime
|
|
|
+from zoneinfo import ZoneInfo
|
|
|
|
|
|
from agents.demand_grade_agent.run import main as grade_demand_words
|
|
|
+from supply_infra.config import get_infra_settings
|
|
|
from supply_infra.db.repositories.demand_grade_repo import DemandGradeRepository
|
|
|
from supply_infra.db.repositories.multi_demand_pool_di_repo import MultiDemandPoolDiRepository
|
|
|
from supply_infra.db.session import get_session
|
|
|
@@ -20,69 +25,107 @@ from supply_infra.db.session import get_session
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_DEFAULT_BATCH_SIZE = 20
|
|
|
-_DEFAULT_MAX_BATCHES = 200
|
|
|
+_DEFAULT_WORKERS = 5
|
|
|
|
|
|
|
|
|
def _resolve_biz_dt(biz_dt: str | None) -> str | None:
|
|
|
if biz_dt:
|
|
|
return biz_dt
|
|
|
- with get_session() as session:
|
|
|
- return MultiDemandPoolDiRepository(session).get_latest_biz_dt()
|
|
|
+ # 定时任务默认按“当天业务日”执行;使用 scheduler_timezone 以免机器时区不同导致 biz_dt 偏移。
|
|
|
+ settings = get_infra_settings()
|
|
|
+ tz = ZoneInfo(settings.scheduler_timezone)
|
|
|
+ return datetime.now(tz).strftime("%Y%m%d")
|
|
|
|
|
|
|
|
|
-def _fetch_next_batch(biz_dt: str, batch_size: int, graded_names: set[str]) -> list[str]:
|
|
|
- """按 exclude_names 过滤已分级词,取下一批(不用 offset,避免和落库进度错位)。"""
|
|
|
+def _fetch_next_batch(biz_dt: str, batch_size: int, exclude_names: set[str]) -> list[str]:
|
|
|
+ """按 exclude_names 过滤已分配/已分级词,取下一批(不用 offset,避免和落库进度错位)。"""
|
|
|
with get_session() as session:
|
|
|
summaries = MultiDemandPoolDiRepository(session).list_distinct_demand_name_summaries(
|
|
|
biz_dt,
|
|
|
limit=batch_size,
|
|
|
offset=0,
|
|
|
- exclude_names=list(graded_names) if graded_names else None,
|
|
|
+ exclude_names=list(exclude_names) if exclude_names else None,
|
|
|
)
|
|
|
return [item["demand_name"] for item in summaries]
|
|
|
|
|
|
|
|
|
+def _allocate_parallel_batches(
|
|
|
+ biz_dt: str,
|
|
|
+ batch_size: int,
|
|
|
+ exclude_names: set[str],
|
|
|
+ *,
|
|
|
+ num_workers: int,
|
|
|
+) -> list[list[str]]:
|
|
|
+ """
|
|
|
+ 为同一轮并行 worker 分配互不重叠的批次。
|
|
|
+
|
|
|
+ 每取出一批后立即加入 reserved,下一批查询时排除,保证单轮内不会重复分配。
|
|
|
+ """
|
|
|
+ batches: list[list[str]] = []
|
|
|
+ reserved = set(exclude_names)
|
|
|
+ for _ in range(num_workers):
|
|
|
+ batch = _fetch_next_batch(biz_dt, batch_size, reserved)
|
|
|
+ if not batch:
|
|
|
+ break
|
|
|
+ batches.append(batch)
|
|
|
+ reserved.update(batch)
|
|
|
+ return batches
|
|
|
+
|
|
|
+
|
|
|
def _fetch_graded_names(biz_dt: str) -> set[str]:
|
|
|
with get_session() as session:
|
|
|
return DemandGradeRepository(session).get_existing_demand_names(biz_dt)
|
|
|
|
|
|
|
|
|
+def _run_batch_in_thread(batch: list[str], biz_dt: str, batch_label: str) -> None:
|
|
|
+ logger.info("Grade demand pool %s (size=%d): %s", batch_label, len(batch), batch)
|
|
|
+ grade_demand_words(batch, biz_dt=biz_dt)
|
|
|
+
|
|
|
+
|
|
|
def grade_demand_pool(
|
|
|
biz_dt: str | None = None,
|
|
|
*,
|
|
|
batch_size: int = _DEFAULT_BATCH_SIZE,
|
|
|
- max_batches: int | None = _DEFAULT_MAX_BATCHES,
|
|
|
+ max_batches: int | None = None,
|
|
|
+ workers: int = _DEFAULT_WORKERS,
|
|
|
) -> dict:
|
|
|
"""
|
|
|
- 循环对需求池中尚未分级的需求打分级,每批调用一次 agent 并等待其完成。
|
|
|
+ 循环对需求池中尚未分级的需求打分级,每轮并行调用多个 agent 并等待其完成。
|
|
|
|
|
|
Args:
|
|
|
biz_dt: 业务日期 YYYYMMDD;不传则取需求池最新业务日。
|
|
|
batch_size: 每批交给 agent 的需求词数量。
|
|
|
- max_batches: 本次运行最多执行多少批(安全阀,避免单次运行时间过长/无限循环);
|
|
|
- None 表示不限制,跑到没有待分级需求为止。
|
|
|
+ max_batches: 本次运行最多执行多少批(安全阀,避免单次运行时间过长/无限循环)。
|
|
|
+ 不传/为 None 时,将根据当天数据库待分级需求词数量动态计算上限:
|
|
|
+ `ceil(total / batch_size) + 10`。
|
|
|
+ workers: 并行线程数;每轮最多同时跑 workers 个互不重叠的批次。
|
|
|
|
|
|
Returns:
|
|
|
运行统计:{"biz_dt", "total", "graded_before", "graded_after", "batches_run", "stopped_reason"}
|
|
|
"""
|
|
|
resolved_biz_dt = _resolve_biz_dt(biz_dt)
|
|
|
- if not resolved_biz_dt:
|
|
|
- result = {"stopped_reason": "no_biz_dt", "batches_run": 0}
|
|
|
- logger.warning("Grade demand pool: no biz_dt available, skip. result=%s", result)
|
|
|
- return result
|
|
|
-
|
|
|
with get_session() as session:
|
|
|
total = MultiDemandPoolDiRepository(session).count_distinct_demand_names(resolved_biz_dt)
|
|
|
|
|
|
graded_names = _fetch_graded_names(resolved_biz_dt)
|
|
|
graded_before = len(graded_names)
|
|
|
|
|
|
+ # max_batches 默认为动态上限:ceil(total / batch_size) + 10
|
|
|
+ # 其中 total 使用“当天需求词去重数”,与本 job 的分页维度一致。
|
|
|
+ if max_batches is None:
|
|
|
+ batches_needed = math.ceil(total / batch_size) if batch_size > 0 else 0
|
|
|
+ max_batches = batches_needed + 10
|
|
|
+
|
|
|
+ worker_count = max(1, workers)
|
|
|
+
|
|
|
logger.info(
|
|
|
- "Grade demand pool start: biz_dt=%s total=%d graded=%d batch_size=%d max_batches=%s",
|
|
|
+ "Grade demand pool start: biz_dt=%s total=%d graded=%d batch_size=%d "
|
|
|
+ "workers=%d max_batches=%s",
|
|
|
resolved_biz_dt,
|
|
|
total,
|
|
|
graded_before,
|
|
|
batch_size,
|
|
|
+ worker_count,
|
|
|
max_batches,
|
|
|
)
|
|
|
|
|
|
@@ -94,34 +137,63 @@ def grade_demand_pool(
|
|
|
stopped_reason = "max_batches_reached"
|
|
|
break
|
|
|
|
|
|
- batch = _fetch_next_batch(resolved_biz_dt, batch_size, graded_names)
|
|
|
- if not batch:
|
|
|
- stopped_reason = "no_more_pending"
|
|
|
+ if max_batches is not None:
|
|
|
+ slots = max_batches - batches_run
|
|
|
+ num_workers = min(worker_count, slots)
|
|
|
+ else:
|
|
|
+ num_workers = worker_count
|
|
|
+
|
|
|
+ if num_workers <= 0:
|
|
|
+ stopped_reason = "max_batches_reached"
|
|
|
break
|
|
|
|
|
|
- batches_run += 1
|
|
|
- logger.info(
|
|
|
- "Grade demand pool batch %d (size=%d): %s",
|
|
|
- batches_run,
|
|
|
- len(batch),
|
|
|
- batch,
|
|
|
+ batches = _allocate_parallel_batches(
|
|
|
+ resolved_biz_dt,
|
|
|
+ batch_size,
|
|
|
+ graded_names,
|
|
|
+ num_workers=num_workers,
|
|
|
)
|
|
|
+ if not batches:
|
|
|
+ stopped_reason = "no_more_pending"
|
|
|
+ break
|
|
|
|
|
|
- try:
|
|
|
- grade_demand_words(batch, biz_dt=resolved_biz_dt)
|
|
|
- except Exception as e:
|
|
|
- logger.error(
|
|
|
- "Grade demand pool batch %d failed: %s", batches_run, e, exc_info=True
|
|
|
- )
|
|
|
+ round_start_batch = batches_run + 1
|
|
|
+ round_failed = False
|
|
|
+
|
|
|
+ with ThreadPoolExecutor(max_workers=len(batches)) as executor:
|
|
|
+ futures = {
|
|
|
+ executor.submit(
|
|
|
+ _run_batch_in_thread,
|
|
|
+ batch,
|
|
|
+ resolved_biz_dt,
|
|
|
+ f"batch {round_start_batch + idx}",
|
|
|
+ ): batch
|
|
|
+ for idx, batch in enumerate(batches)
|
|
|
+ }
|
|
|
+ for future in as_completed(futures):
|
|
|
+ batch = futures[future]
|
|
|
+ try:
|
|
|
+ future.result()
|
|
|
+ except Exception as e:
|
|
|
+ logger.error(
|
|
|
+ "Grade demand pool batch failed: %s (batch=%s)",
|
|
|
+ e,
|
|
|
+ batch,
|
|
|
+ exc_info=True,
|
|
|
+ )
|
|
|
+ round_failed = True
|
|
|
+
|
|
|
+ batches_run += len(batches)
|
|
|
+
|
|
|
+ if round_failed:
|
|
|
stopped_reason = "batch_failed"
|
|
|
break
|
|
|
|
|
|
new_graded_names = _fetch_graded_names(resolved_biz_dt)
|
|
|
if len(new_graded_names) <= len(graded_names):
|
|
|
- # agent 没有对这批产生任何新的落库结果,避免死循环重复拿到同一批
|
|
|
+ # agent 没有对本轮产生任何新的落库结果,避免死循环重复拿到同一批
|
|
|
logger.warning(
|
|
|
- "Grade demand pool batch %d made no progress (graded still %d), stop",
|
|
|
- batches_run,
|
|
|
+ "Grade demand pool round made no progress (graded still %d), stop",
|
|
|
len(graded_names),
|
|
|
)
|
|
|
stopped_reason = "stalled"
|
|
|
@@ -135,6 +207,7 @@ def grade_demand_pool(
|
|
|
"graded_before": graded_before,
|
|
|
"graded_after": len(graded_names),
|
|
|
"batches_run": batches_run,
|
|
|
+ "workers": worker_count,
|
|
|
"stopped_reason": stopped_reason,
|
|
|
"run_at": datetime.now().isoformat(),
|
|
|
}
|