Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions crates/utopia-server/src/api/admin_routes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ pub async fn get_deployment(
#[derive(Deserialize)]
pub struct DeploymentReq {
pub open_registration: bool,
/// 任务 worker 并发(1-256):**外层兜底**,防任务无限堆积。
/// 真正的节流是按模型的限额,这个值应当明显大于各模型限额之和
/// 任务 worker 并发(1-256):管模型任务,另有 `jobs::NO_CHAT_LANE` 个名额留给入库任务。
/// 它是**外层兜底**,防任务无限堆积;真正的节流是按模型的限额,这个值应当明显大于各模型限额之和
#[serde(default)]
pub worker_concurrency: Option<i32>,
/// 未单独配置的模型走的并发缺省
Expand Down
39 changes: 30 additions & 9 deletions crates/utopia-store/src/jobs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -373,20 +373,27 @@ pub async fn status_in_kb(pool: &PgPool, kb_id: Uuid, id: i64) -> AppResult<Opti
/// 认领一个到期任务;没有则返回 None。
///
/// 不调对话模型的任务排在前面(`job_priority`,迁移 0116,#1300):同步、文档入库这些几秒
/// 的活,不该等在几百个更早的抽取后面
async fn claim_one(pool: &PgPool) -> AppResult<Option<Job>> {
let job = sqlx::query_as(
/// 的活,不该等在几百个更早的抽取后面。`no_chat_only` 只认领这些,给保留通道用
async fn claim_one(pool: &PgPool, no_chat_only: bool) -> AppResult<Option<Job>> {
// 写成字面量而不是绑定参数:这样它是索引条件,扫到优先级 0 的尽头就停;绑定参数只能
// 逐行筛,模型任务堆到上万条时,常规名额占满后每 200ms 都要把它们全扫一遍
let only = if no_chat_only {
"AND job_priority(kind) = 0"
} else {
""
};
let job = sqlx::query_as(&format!(
"UPDATE jobs SET status = 'running', locked_at = now(),
attempts = attempts + 1, updated_at = now()
WHERE id = (
SELECT id FROM jobs
WHERE status = 'queued' AND run_at <= now()
WHERE status = 'queued' AND run_at <= now() {only}
ORDER BY job_priority(kind), run_at
FOR UPDATE SKIP LOCKED
LIMIT 1
)
RETURNING id, kind, payload, attempts, max_attempts, locked_at",
)
RETURNING id, kind, payload, attempts, max_attempts, locked_at"
))
.fetch_optional(pool)
.await?;
Ok(job)
Expand Down Expand Up @@ -586,8 +593,16 @@ async fn persist_with_retry(
}
}

/// 常规名额全占满之后,还能再认领几个不调对话模型的任务(#1300)。
///
/// 光排序不够:抽取等模型许可的时候照样占着 worker 名额。本地模型上名额常被等许可的
/// 抽取占满,排在最前面的同步也只能等某个抽取跑完。这两个名额只给 `job_priority` 为 0
/// 的任务,所以 `worker_concurrency` 管的是其余一切,同时在跑的最多再多两个
pub const NO_CHAT_LANE: usize = 2;

/// worker 调度循环:运行中任务数低于目标并发就继续认领(有活立即续派),
/// 空闲时等入队通知、2s 轮询兜底(#517);每个任务在独立 tokio task 中执行,长抽取不再阻塞同步。
/// 常规名额满了,还有 [`NO_CHAT_LANE`] 个名额只认领不调对话模型的任务。
/// `concurrency` 每轮热读——系统设置里改并发数即时生效。
/// 任务分发逻辑由调用方以 handler 注入(store 不依赖上层 crate)。
pub async fn run_worker<F, Fut>(pool: PgPool, concurrency: Arc<AtomicUsize>, handler: F)
Expand All @@ -596,6 +611,7 @@ where
Fut: std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
{
let running = Arc::new(AtomicUsize::new(0));
let in_lane = Arc::new(AtomicUsize::new(0));
// 孤儿回收:进程被杀时 running 任务无人收尸,文档会永远停在 extracting。
// 单进程部署下,启动时仍为 running 的必是孤儿——一律重排队(处理器均幂等)。
match sqlx::query(
Expand Down Expand Up @@ -624,16 +640,19 @@ where
return;
}
let cap = concurrency.load(Ordering::Relaxed).max(1);
if running.load(Ordering::Relaxed) >= cap {
// 常规名额满了就只往保留通道里认领,通道也满了才歇
let lane = running.load(Ordering::Relaxed) >= cap;
if lane && in_lane.load(Ordering::Relaxed) >= NO_CHAT_LANE {
tokio::time::sleep(Duration::from_millis(200)).await;
continue;
}
match claim_one(&pool).await {
match claim_one(&pool, lane).await {
Ok(Some(job)) => {
// 从哪边的名额认领的,跑完就还给哪边
let running = if lane { &in_lane } else { &running }.clone();
running.fetch_add(1, Ordering::Relaxed);
let pool = pool.clone();
let handler = handler.clone();
let running = running.clone();
tokio::spawn(async move {
// **处理器 panic 也要收尸。** 直接 `handler(job).await` 的话,
// panic 会把这个 spawn 出来的 future 一起掀掉:`mark_failed`
Expand Down Expand Up @@ -663,6 +682,8 @@ where
running.fetch_sub(1, Ordering::Relaxed);
});
}
// 常规名额满着:随时可能空出一个,照旧 200ms 再看
Ok(None) if lane => tokio::time::sleep(Duration::from_millis(200)).await,
Ok(None) => {
wait_for_work(listener.as_mut(), IDLE_POLL).await;
}
Expand Down
2 changes: 1 addition & 1 deletion crates/utopia-store/src/jobs_claim_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ async fn a_later_sync_is_claimed_before_an_earlier_extraction() -> anyhow::Resul
.fetch_all(&pool)
.await?;

let claimed = claim_one(&pool).await?.map(|job| job.kind);
let claimed = claim_one(&pool, false).await?.map(|job| job.kind);
sqlx::query("DELETE FROM jobs WHERE id = ANY($1)")
.bind(&ids)
.execute(&pool)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
//! #1300:常规名额全被调模型的任务占着时,后来的同步照样有名额跑(`jobs::NO_CHAT_LANE`)。
//!
//! 起真实的 `run_worker`,它会认领库里任何排队任务,所以自成一个二进制(见
//! `tests/store/main.rs`)。处理器只认带本测试标记的两条,别人的任务一概报错,不替人干活。

use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::time::Duration;
use utopia_store::jobs;

#[tokio::test]
async fn a_sync_runs_while_model_work_holds_every_slot() -> anyhow::Result<()> {
let Some(url) = utopia_store::test_db::url() else {
return Ok(());
};
let pool = sqlx::PgPool::connect(&url).await?;
let mark = uuid::Uuid::now_v7().to_string();
let (started, mut extracting) = tokio::sync::mpsc::unbounded_channel();
let (finished, mut ran) = tokio::sync::mpsc::unbounded_channel();
let ours = mark.clone();
let worker = tokio::spawn(jobs::run_worker(
pool.clone(),
Arc::new(AtomicUsize::new(1)),
move |job| {
let (ours, started, finished) = (ours.clone(), started.clone(), finished.clone());
async move {
anyhow::ensure!(
job.payload["lane_test"] == ours.as_str(),
"not this test's job"
);
if job.kind == "extract_document" {
started.send(())?;
// 一直占着唯一的常规名额,像一次在等模型许可的抽取
std::future::pending::<()>().await;
}
finished.send(job.kind)?;
Ok(())
}
},
));
let queue = |kind: &'static str| {
// 放在一百年前,排在同优先级里库里别的一切前面
sqlx::query("INSERT INTO jobs (kind, payload, run_at) VALUES ($1, $2, now() - interval '100 years')")
.bind(kind)
.bind(serde_json::json!({ "lane_test": mark }))
.execute(&pool)
};
let result = async {
queue("extract_document").await?;
tokio::time::timeout(Duration::from_secs(10), extracting.recv()).await?;
queue("sync_source").await?;
let synced = tokio::time::timeout(Duration::from_secs(10), ran.recv()).await;
anyhow::ensure!(
matches!(&synced, Ok(Some(kind)) if kind == "sync_source"),
"the sync waited behind the extraction: {synced:?}"
);
anyhow::Ok(())
}
.await;
worker.abort();
let _ = worker.await;
sqlx::query("DELETE FROM jobs WHERE payload->>'lane_test' = $1")
.bind(&mark)
.execute(&pool)
.await?;
result
}
2 changes: 2 additions & 0 deletions crates/utopia-store/tests/store/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
//! - `a_due_base_is_queued_once`:`queue_due_inferences` 扫全库到期的库。0050 起新建的库开着
//! 推理又从没推过,别的测试刚建的库全算到期,会被它排上推理任务——数自己库任务的测试
//! (比如 lib 单元测试里的 `phrase_bindings_delivery_tests`)就会多数出一个。
//! - `a_sync_runs_while_model_work_holds_every_slot`:同样起真实的 `run_worker`,同样会认领
//! 库里任何排队任务,不是自己的就报错。
//!
//! 新加测试默认放这个目录;只有碰上面那类全局状态时才放顶层,并在这里补一行为什么。
mod a_batch_decides_like_a_person;
Expand Down
8 changes: 5 additions & 3 deletions web/src/i18n/en.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1194,9 +1194,11 @@ export const en = {
"When off, the sign-up form is closed and only admins can create accounts here.",
workers: "Background workers",
workersHint:
"An outer ceiling on how many jobs run at once (1–256), there to stop work piling up " +
"without bound. The real throttle is the per-model limit below, so keep this comfortably " +
"above the sum of those. Takes effect immediately.",
"How many model jobs (extraction, resolution) can run at once (1–256), plus two slots " +
"reserved for intake jobs such as syncs, which make no model call and so never wait " +
"behind model work. It is an outer ceiling that stops work piling up without bound; the " +
"real throttle is the per-model limit below, so keep this comfortably above the sum of " +
"those. Takes effect immediately.",
workersApply: "Apply",
/* 真正的节流:约束来自供应商的速率限制,而那是按模型算的 */
modelConcurrency: "Model concurrency",
Expand Down
3 changes: 2 additions & 1 deletion web/src/i18n/zh.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1057,7 +1057,8 @@ export const zh: Strings = {
openRegHint: "关闭后注册表单不再开放,只有管理员能在此创建账户。",
workers: "后台工作线程",
workersHint:
"同时运行的任务数的外层上限(1–256),作用是防止工作无限堆积。" +
"最多同时运行多少个模型任务(抽取、消解,1–256),另有两个名额保留给同步这类入库任务:" +
"它们不调模型,所以不必排在模型任务后面。这是防止工作无限堆积的外层上限;" +
"真正的节流是下面的按模型上限,所以这个数留得比它们之和宽裕些。即时生效。",
workersApply: "应用",
modelConcurrency: "模型并发",
Expand Down
Loading