diff --git a/cloud/cost.py b/cloud/cost.py index b052830..70cb869 100644 --- a/cloud/cost.py +++ b/cloud/cost.py @@ -76,6 +76,41 @@ class CostStore: ) self._conn.commit() + def _new_conn(self): + """每次写/读开一个新连接，用完即关。 + + 旧实现在 __init__ 里建一条连接并终身共用，单线程没问题； + worker 并发后多个线程共用同一条 sqlite 连接会直接报 + ProgrammingError（连接是线程绑定的）。其它 store（Worker + 的 _set_status、_audit 等）本来就是每次开新连接，这里 + 向它们看齐。_conn 仍保留（建表 + 旧测试直接用它读）， + 但 record/_report 不再走它。 + 例外：db_path 为 :memory: 时新连接是另一个空库，必须 + 复用 _conn（:memory: 只用于单线程单测，并发 worker + 从不用它）。 + """ + if self.db_path == ":memory:": + return self._conn + conn = _db_connect(self.db_path) + conn.execute( + """CREATE TABLE IF NOT EXISTS cost_records ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + task_id TEXT, + project_id TEXT, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + total_tokens INTEGER, + usd_micros INTEGER, + source TEXT, + fallback_count INTEGER DEFAULT 0, + day TEXT, + created_at REAL + )""" + ) + conn.commit() + return conn + def record(self, task_id, project_id, model, input_tokens=0, output_tokens=0, usd_micros=None, source="actual", fallback_count=0): @@ -93,18 +128,24 @@ class CostStore: total_tokens = input_tokens + output_tokens created_at = time.time() day = datetime.fromtimestamp(created_at, tz=timezone.utc).strftime("%Y-%m-%d") - cursor = self._conn.execute( - """INSERT INTO cost_records - (task_id, project_id, model, input_tokens, output_tokens, - total_tokens, usd_micros, source, fallback_count, day, - created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - (task_id, project_id, model, input_tokens, output_tokens, - total_tokens, usd_micros, source, int(fallback_count or 0), - day, created_at), - ) - self._conn.commit() - return cursor.lastrowid + conn = self._new_conn() + try: + cursor = conn.execute( + """INSERT INTO cost_records + (task_id, project_id, model, input_tokens, output_tokens, + total_tokens, usd_micros, source, fallback_count, day, + created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (task_id, project_id, model, input_tokens, output_tokens, + total_tokens, usd_micros, source, int(fallback_count or 0), + day, created_at), + ) + conn.commit() + return cursor.lastrowid + finally: + # :memory: 时 conn 就是共享的 _conn，绝不能关掉它。 + if conn is not self._conn: + conn.close() def record_estimated(self, task_id, project_id, model, prompt_text, output_text): @@ -123,19 +164,24 @@ class CostStore: ) def _report(self, where, params): - row = self._conn.execute( - """SELECT - COALESCE(SUM(total_tokens), 0), - COALESCE(SUM(usd_micros), 0), - COUNT(*), - COALESCE(SUM(fallback_count), 0), - COALESCE(SUM(CASE WHEN source = 'actual' - THEN total_tokens ELSE 0 END), 0), - COALESCE(SUM(CASE WHEN source = 'estimated' - THEN total_tokens ELSE 0 END), 0) - FROM cost_records WHERE %s""" % where, - params, - ).fetchone() + conn = self._new_conn() + try: + row = conn.execute( + """SELECT + COALESCE(SUM(total_tokens), 0), + COALESCE(SUM(usd_micros), 0), + COUNT(*), + COALESCE(SUM(fallback_count), 0), + COALESCE(SUM(CASE WHEN source = 'actual' + THEN total_tokens ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN source = 'estimated' + THEN total_tokens ELSE 0 END), 0) + FROM cost_records WHERE %s""" % where, + params, + ).fetchone() + finally: + if conn is not self._conn: + conn.close() total_tokens, usd_micros, turns, fallback_count, actual, estimated = row return { "total_tokens": int(total_tokens), diff --git a/cloud/sandbox.py b/cloud/sandbox.py index 6626fd4..fe6b41f 100644 --- a/cloud/sandbox.py +++ b/cloud/sandbox.py @@ -420,10 +420,17 @@ class DockerExecutor(Executor): return cmd # -- runner 调用 / 清理 -------------------------------------------- - def _call_runner(self, cmd): - """调 runner：优先传 timeout，假 runner 只接受 cmd 时回退。""" + def _call_runner(self, cmd, timeout=None): + """调 runner：优先传 timeout，假 runner 只接受 cmd 时回退。 + + timeout 参数只影响本次调用（并发下两个任务各用各的 + 超时）；不传时用实例的 self.timeout。绝不能在调用方 + 临时改 self.timeout 再恢复，那在并发下会互相覆盖。 + """ try: - return self.runner(cmd, timeout=self.timeout) + return self.runner( + cmd, + timeout=self.timeout if timeout is None else timeout) except TypeError: return self.runner(cmd) @@ -493,20 +500,20 @@ class DockerExecutor(Executor): continue return ok - def _run_in_container(self, task, inner_cmd): + def _run_in_container(self, task, inner_cmd, timeout=None): """在容器里跑一轮内层命令，返回 LocalExecutor 同款结果 dict。 超时：runner 抛 subprocess.TimeoutExpired（或同名异常） 时先 kill+rm 容器，再返回 failed（stdout 截断 + timeout 说明，returncode 为 None）；其它 runner 异常也如实返回 - failed，不编造成功。 + failed，不编造成功。timeout 只影响本次调用，不改实例。 """ task_id = task.get("id") if isinstance(task, dict) else None if task_id is None: raise ValueError("任务缺少 id") cmd = self.build_run_command(task, inner_cmd) try: - result = self._call_runner(cmd) + result = self._call_runner(cmd, timeout=timeout) except Exception as exc: if exc.__class__.__name__ == "TimeoutExpired": self.kill_container(task) @@ -529,7 +536,7 @@ class DockerExecutor(Executor): parsed["container"] = self.container_name(task) return parsed - def dispatch_task(self, task) -> dict: + def dispatch_task(self, task, timeout=None) -> dict: """容器内调 host.py dispatch：建 host 任务并跑一轮 turn。""" if not isinstance(task, dict) or task.get("id") is None: raise ValueError("任务缺少 id") @@ -541,9 +548,9 @@ class DockerExecutor(Executor): repo = task.get("repo") if repo: inner += ["--repo", repo] - return self._run_in_container(task, inner) + return self._run_in_container(task, inner, timeout=timeout) - def run_task(self, task) -> dict: + def run_task(self, task, timeout=None) -> dict: """容器内调 host.py run：在已有 host 任务上再跑一轮 turn。""" if not isinstance(task, dict) or task.get("id") is None: raise ValueError("任务缺少 id") @@ -552,7 +559,7 @@ class DockerExecutor(Executor): inner = self._inner_host_command( "run", "--task", self._host_name(task), "--prompt", task["prompt"], model=task.get("model")) - return self._run_in_container(task, inner) + return self._run_in_container(task, inner, timeout=timeout) def execute(self, task) -> dict: """只读执行：容器内调 host.py status，不直接跑模型回合。 diff --git a/cloud/server.py b/cloud/server.py index b669f39..83e83ae 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -624,6 +624,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): self._merge_task(user_id, rest[: -len("/merge")]) elif rest.endswith("/pr") and method == "POST": self._task_pr(user_id, rest[: -len("/pr")]) + elif rest.endswith("/retry") and method == "POST": + self._retry_task(user_id, rest[: -len("/retry")]) elif rest.endswith("/heartbeat") and method == "POST": self._heartbeat_task(user_id, rest[: -len("/heartbeat")]) else: @@ -1321,6 +1323,67 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): store.audit(user_id, task["project_id"], "merge_task", task_id) self._send_json(200, task) + def _retry_task(self, user_id, task_id): + """把终态失败任务重新排队：只改状态，不直接执行。 + + 规则：任务必须存在（404）、调用方必须是项目 maintainer + 以上（approve_merge，403——重跑要再花模型钱，不能让 + 只有读权限或普通贡献者随便点）、任务必须是 failed / + interrupted / merge-failed（其它状态一律 409：running + 本来就在跑，queued 本来就在等，completed/merged 已有 + 结果，重跑应是新任务而不是覆盖结果）。 + 成功后状态回 queued，上一轮的 result_text/error_text/ + actual_model 清空（否则新一轮的结果会和旧结果混在一起 + 无法分辨）；prompt、模型、优先级、依赖原样保留，依赖 + 已完成时重跑任务会立刻就绪。写审计 task_retried。 + 更新用 CAS（WHERE status IN (...)），避免与正在完成的 + 回合互相覆盖：若 CAS 没改到行，说明状态已变，返回 + 409 而不是假装成功。 + """ + row = self._fetch_task(task_id) + if row is None: + self._send_json(404, {"error": "task not found"}) + return + task = _task_to_dict(row) + with _auth_store(self.server.db_path) as store: + if not store.check_permission( + user_id, task["project_id"], "approve_merge"): + self._send_json(403, {"error": "forbidden"}) + return + if task["status"] not in ("failed", "interrupted", "merge-failed"): + self._send_json(409, { + "error": "only failed/interrupted tasks can be retried", + "status": task["status"], + }) + return + now = time.time() + conn = self._db() + try: + cur = conn.execute( + "UPDATE tasks SET status = 'queued', result_text = NULL," + " error_text = NULL, actual_model = NULL, updated_at = ?" + " WHERE id = ? AND status IN" + " ('failed', 'interrupted', 'merge-failed')", + (now, task_id), + ) + conn.commit() + if cur.rowcount != 1: + self._send_json(409, { + "error": "task status changed, retry aborted", + "status": task["status"], + }) + return + task = _task_to_dict( + conn.execute( + "SELECT * FROM tasks WHERE id = ?", (task_id,) + ).fetchone() + ) + finally: + conn.close() + with _auth_store(self.server.db_path) as store: + store.audit(user_id, task["project_id"], "task_retried", task_id) + self._send_json(200, task) + # -- worker 池 ----------------------------------------------------- def _workers_list(self, user_id): """列出注册 worker：任何已认证用户可读（worker 是基础设施 diff --git a/cloud/tests/test_docker_sandbox.py b/cloud/tests/test_docker_sandbox.py index 2470a71..d02b26f 100644 --- a/cloud/tests/test_docker_sandbox.py +++ b/cloud/tests/test_docker_sandbox.py @@ -289,8 +289,7 @@ class WorkerDockerTestCase(unittest.TestCase): worker = Worker(self.db_path, "/fake/host.py", env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex) # 模拟 dispatch 阻塞期间：worker 已把容器名记下来。 - worker._running_task_id = "d3" - worker._running_container = "codex-d3" + worker._running_set("d3", container="codex-d3") conn = sqlite3.connect(self.db_path) try: conn.execute( diff --git a/cloud/tests/test_server_retry.py b/cloud/tests/test_server_retry.py new file mode 100644 index 0000000..02a2f63 --- /dev/null +++ b/cloud/tests/test_server_retry.py @@ -0,0 +1,136 @@ +"""失败任务重跑端点测试：POST /tasks/{id}/retry。 + +只测状态机本身，不执行真实任务：失败/中断任务重回 queued， +上一轮结果清空，prompt 与模型原样保留；权限与状态码语义 +与 merge 端点对齐（maintainer 以上、404/401/403/409）。 +""" + +import http.client +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import unittest +from pathlib import Path + +try: + from cloud.auth import AuthStore + from cloud.server import create_server +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.auth import AuthStore + from cloud.server import create_server + + +class RetryTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + self.auth = AuthStore(self.db_path) + self.owner_id, self.owner_key = self.auth.create_user("owner") + self.contrib_id, self.contrib_key = self.auth.create_user("contrib") + self.server = create_server(self.db_path, "127.0.0.1", 0) + self.port = self.server.server_address[1] + self.thread = threading.Thread( + target=self.server.serve_forever, daemon=True) + self.thread.start() + status, self.project = self._req( + "POST", "/projects", {"name": "demo"}, self.owner_key) + self.assertEqual(status, 201) + self.auth.set_member( + self.project["id"], self.owner_id, self.contrib_id, "contributor") + + def tearDown(self): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + try: + self.auth.conn.close() + except Exception: + pass + self.tmp.cleanup() + + def _req(self, method, path, body=None, key=None): + conn = http.client.HTTPConnection( + "127.0.0.1", self.port, timeout=10) + hdrs = {} + payload = None + if body is not None: + payload = json.dumps(body).encode("utf-8") + hdrs["Content-Type"] = "application/json" + if key is not None: + hdrs["X-API-Key"] = key + try: + conn.request(method, path, body=payload, headers=hdrs) + resp = conn.getresponse() + data = resp.read() + finally: + conn.close() + try: + parsed = json.loads(data.decode("utf-8")) if data else None + except ValueError: + parsed = None + return resp.status, parsed + + def _failed_task(self): + status, task = self._req( + "POST", "/tasks", + {"project_id": self.project["id"], "prompt": "x"}, + self.owner_key) + self.assertEqual(status, 201) + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "UPDATE tasks SET status = 'failed'," + " error_text = 'old boom', actual_model = 'm1'" + " WHERE id = ?", + (task["id"],)) + conn.commit() + finally: + conn.close() + return task + + def test_retry_failed_by_owner_returns_queued_and_clears_result(self): + task = self._failed_task() + status, body = self._req( + "POST", "/tasks/%s/retry" % task["id"], {}, self.owner_key) + self.assertEqual(status, 200) + self.assertEqual(body["status"], "queued") + self.assertIsNone(body.get("error_text")) + self.assertIsNone(body.get("actual_model")) + conn = sqlite3.connect(self.db_path) + try: + actions = [r[0] for r in conn.execute( + "SELECT action FROM audit_log WHERE project_id = ?", + (self.project["id"],)).fetchall()] + finally: + conn.close() + self.assertIn("task_retried", actions) + + def test_retry_non_failed_409(self): + status, task = self._req( + "POST", "/tasks", + {"project_id": self.project["id"], "prompt": "x"}, + self.owner_key) + self.assertEqual(status, 201) + status, body = self._req( + "POST", "/tasks/%s/retry" % task["id"], {}, self.owner_key) + self.assertEqual(status, 409) + + def test_retry_by_contributor_forbidden_and_missing_404(self): + task = self._failed_task() + status, _ = self._req( + "POST", "/tasks/%s/retry" % task["id"], {}, self.contrib_key) + self.assertEqual(status, 403) + status, _ = self._req( + "POST", "/tasks/nope/retry", {}, self.owner_key) + self.assertEqual(status, 404) + status, _ = self._req( + "POST", "/tasks/%s/retry" % task["id"], {}) + self.assertEqual(status, 401) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/tests/test_worker_concurrency.py b/cloud/tests/test_worker_concurrency.py new file mode 100644 index 0000000..0601b18 --- /dev/null +++ b/cloud/tests/test_worker_concurrency.py @@ -0,0 +1,188 @@ +"""Worker 并发单测。 + +只用注入的假 runner（假 CLI），不跑真实模型、不联网。 +capacity>1 时同一 worker 内多个任务必须真正并行跑，而不是 +被单任务字段或共享连接拖回串行；capacity=1 时行为与旧串行 +完全一致；最终状态写回必须是 CAS，取消/回收后的任务不能 +被后台线程盖回 completed。 +""" + +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import time +import unittest +from pathlib import Path +from types import SimpleNamespace + +try: + from cloud.server import init_db + from cloud.worker import Worker +except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.server import init_db + from cloud.worker import Worker + + +def _insert_task(db_path, task_id, prompt="do work", project_id="p1"): + now = time.time() + conn = sqlite3.connect(db_path) + try: + conn.execute( + "INSERT INTO tasks(id, project_id, prompt, status," + " idempotency_key, depends_on, priority, created_by," + " created_at, updated_at, last_heartbeat)" + " VALUES(?, ?, ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)", + (task_id, project_id, prompt, now, now), + ) + conn.commit() + finally: + conn.close() + + +def _status(db_path, task_id): + conn = sqlite3.connect(db_path) + try: + row = conn.execute( + "SELECT status FROM tasks WHERE id = ?", + (task_id,)).fetchone() + finally: + conn.close() + return row[0] if row else None + + +def _host_json(name, status, text): + return SimpleNamespace( + returncode=0, + stdout=json.dumps({"task": name, "status": status, "text": text}), + stderr="") + + +class WorkerConcurrencyTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + init_db(self.db_path) + self.env = {"XCM_API_KEY": "test-placeholder"} + + def tearDown(self): + self.tmp.cleanup() + + def _worker(self, runner=None, **kwargs): + kwargs.setdefault("env", self.env) + kwargs.setdefault("retry_backoff_seconds", 0) + kwargs.setdefault("max_retries", 0) + return Worker(self.db_path, "/fake/host.py", + runner=runner, **kwargs) + + def test_capacity_2_runs_two_tasks_concurrently(self): + # 两个任务各睡 0.5 秒：并行时它们必须在同一时间段内 + # 都在跑（用 Barrier 证明重叠，而不是只看总时长，避免 + # 慢 CI 上的时间断言抖动）。串行时第一个 runner 会在 + # Barrier 上等满 5 秒超时，测试会因此失败——这正是要 + # 证明的：capacity=2 不是串行。 + barrier = threading.Barrier(2, timeout=5) + + def fake(cmd): + barrier.wait() + time.sleep(0.1) + return _host_json("x", "completed", "done") + + _insert_task(self.db_path, "A") + _insert_task(self.db_path, "B") + worker = self._worker(runner=fake, capacity=2) + results = worker.run_all() + self.assertEqual(len(results), 2) + self.assertEqual({r["status"] for r in results}, {"completed"}) + self.assertEqual(_status(self.db_path, "A"), "completed") + self.assertEqual(_status(self.db_path, "B"), "completed") + + def test_capacity_1_stays_serial(self): + # capacity=1 时 run_all 仍是旧串行循环：两个任务都完成， + # 且第二个任务的认领发生在第一个完成之后（靠任务状态 + # 变化顺序断言，不靠时间）。 + order = [] + lock = threading.Lock() + + def fake(cmd): + with lock: + order.append("start") + time.sleep(0.2) + with lock: + order.append("end") + return _host_json("x", "completed", "done") + + _insert_task(self.db_path, "A") + _insert_task(self.db_path, "B") + worker = self._worker(runner=fake, capacity=1) + results = worker.run_all() + self.assertEqual(len(results), 2) + self.assertEqual(order, ["start", "end", "start", "end"]) + + def test_pool_max_tasks_respected(self): + # 并行模式下 max_tasks=1 只跑一个任务：提交数受 + # max_tasks 约束，不能因为线程池先提交了两个就多跑。 + def fake(cmd): + time.sleep(0.1) + return _host_json("x", "completed", "done") + + for tid in ("A", "B", "C"): + _insert_task(self.db_path, tid) + worker = self._worker(runner=fake, capacity=2) + results = worker.run_all(max_tasks=1) + self.assertEqual(len(results), 1) + done = [tid for tid in ("A", "B", "C") + if _status(self.db_path, tid) == "completed"] + self.assertEqual(len(done), 1) + + def test_wait_for_running(self): + # _wait_for_running：有任务在跑时超时返回 False， + # 任务清掉后返回 True（run_loop 优雅排空靠它）。 + worker = self._worker(capacity=2) + worker._running_set("A", container="codex-A") + self.assertFalse(worker._wait_for_running(timeout=0.3)) + + def clearer(): + time.sleep(0.2) + worker._running_clear("A") + + threading.Thread(target=clearer, daemon=True).start() + self.assertTrue(worker._wait_for_running(timeout=5)) + + def test_final_write_does_not_overwrite_interrupted(self): + # 任务在跑时被外部标成 interrupted（取消/回收），随后 + # runner 返回 completed，最终写回绝不能把它盖回 completed。 + started = threading.Event() + + def fake(cmd): + started.set() + time.sleep(0.5) + return _host_json("x", "completed", "done") + + _insert_task(self.db_path, "A") + worker = self._worker(runner=fake, capacity=1) + outcome = {} + + def run(): + outcome["result"] = worker.run_one() + + t = threading.Thread(target=run, daemon=True) + t.start() + self.assertTrue(started.wait(timeout=5)) + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "UPDATE tasks SET status = 'interrupted' WHERE id = 'A'") + conn.commit() + finally: + conn.close() + t.join(timeout=10) + self.assertEqual(outcome["result"]["status"], "interrupted") + self.assertEqual(_status(self.db_path, "A"), "interrupted") + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 9f977ed..afef742 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -117,6 +117,7 @@ import signal import socket import subprocess import sys +import threading import time from types import SimpleNamespace @@ -270,13 +271,21 @@ class Worker: repo=repo) self._cost_store = cost_store self._cancel_requested = set() - # 真取消用：正在跑的子进程（Popen 对象）与任务 id。 - # process_factory 可注入：单测用它返回一个假长跑子进程， - # 真实执行时走 LocalExecutor.popen_task（进程组启动）。 + # 真取消用：process_factory 可注入（单测用它返回一个假长跑 + # 子进程，真实执行时走 LocalExecutor.popen_task 进程组启动）。 self.process_factory = process_factory - self._running_proc = None - self._running_task_id = None - self._running_container = None + # 并发运行状态：task_id -> {"proc": Popen|None, + # "container": 容器名|None}，受 _running_lock 保护。 + # 旧版只有单任务三个字段（_running_proc 等），capacity=1 + # 时够用；capacity>1 后多个任务同时跑，单字段会被后 + # 一个任务覆盖，导致 cancel 杀错、退出时漏杀。所有读写 + # 都走下面的 _running_set/_running_clear/_running_snapshot。 + self._running = {} + self._running_lock = threading.Lock() + # DB 写锁：cost 记录共用一条连接（CostStore._conn）， + # usage 计数也是共享变量；模型回合本身在子进程/容器里 + # 跑不受这把锁影响，锁只覆盖毫秒级的写库操作。 + self._db_lock = threading.RLock() self._loop_stop = False self.logger = get_logger("cloud.worker") # 任务级超时 / 失败重试（见模块 docstring 的生产级说明）。 @@ -469,6 +478,27 @@ class Worker: except Exception: pass + def _set_status_if_running(self, task_id, status): + """仅当任务仍为 running 时改状态，返回是否改成功。 + + 并发 + 优雅排空下，后台线程可能在任务已被取消 + （标 interrupted）后才跑完；此时最终写回绝不能把 + interrupted 盖回 completed/failed。正常流程里任务 + 到最终写回时仍是 running，CAS 一定成功。 + """ + now = time.time() + conn = _db_connect(self.db_path) + try: + cur = conn.execute( + "UPDATE tasks SET status = ?, last_heartbeat = ?," + " updated_at = ? WHERE id = ? AND status = 'running'", + (status, now, now, task_id), + ) + conn.commit() + return cur.rowcount == 1 + finally: + conn.close() + def _write_execution_result(self, task_id, result_text=None, error_text=None, actual_model=None, started_at=None, finished_at=None, @@ -597,6 +627,53 @@ class Worker: if t.get("executor") in (None, "", mine)] return ready[0] if ready else None + # -- 并发运行状态 -------------------------------------------------- + def _running_set(self, task_id, proc=None, container=None): + with self._running_lock: + entry = self._running.setdefault(task_id, {}) + if proc is not None: + entry["proc"] = proc + if container is not None: + entry["container"] = container + + def _running_clear(self, task_id): + with self._running_lock: + self._running.pop(task_id, None) + + def _running_snapshot(self): + with self._running_lock: + return {tid: dict(entry) + for tid, entry in self._running.items()} + + def _wait_for_running(self, timeout=None): + """等 _running 变空；timeout 为 None 时无限等。 + + 返回 True 表示已空，False 表示超时仍有任务在跑。 + run_loop 退出时的优雅排空用它：先等任务自然结束， + 超时后调用方再杀残留并标 interrupted。 + """ + deadline = None if timeout is None else time.time() + timeout + while True: + with self._running_lock: + if not self._running: + return True + if deadline is not None and time.time() >= deadline: + return False + time.sleep(0.2) + + def _heartbeat_running_tasks(self): + """给所有正在跑的任务刷 last_heartbeat。 + + run_loop 的心跳线程调它：任务在跑时心跳保持新鲜， + 自动陈旧回收（sweep）才不会误杀活任务；worker 真死后 + 心跳停了，任务超过阈值才会被别的 worker 回收。 + """ + for task_id in self._running_snapshot(): + try: + self.heartbeat(task_id) + except Exception: + pass + # -- 取消（server cancel 语义，见模块 docstring）------------------- def cancel(self, task_id): """把 queued/running 任务标 interrupted；已终态或不存在返回 False。 @@ -608,26 +685,24 @@ class Worker: if task is None or task["status"] in TERMINAL_STATUSES: return False self._cancel_requested.add(task_id) - if self._running_task_id == task_id and self._running_proc is not None: - self._kill_proc_group(self._running_proc) + entry = self._running_snapshot().get(task_id) or {} + if entry.get("proc") is not None: + self._kill_proc_group(entry["proc"]) try: - self._running_proc.wait(timeout=3) + entry["proc"].wait(timeout=3) except Exception: pass - self._running_proc = None - self._running_task_id = None # Docker 取消：按任务派生的容器名真杀，不只改数据库状态。 - # 同步路径（runner 阻塞在 docker run 里）下 _running_proc - # 为空，此时靠 _running_container 杀容器；kill 尽力而为， - # 失败也不影响后续的状态标记。 - if self._running_task_id == task_id and \ - self._running_container is not None and \ + # 同步路径（runner 阻塞在 docker run 里）下 proc 为空， + # 此时靠容器名杀容器；kill 尽力而为，失败也不影响后续 + # 的状态标记。 + if entry.get("container") is not None and \ hasattr(self.executor, "kill_container"): try: - self.executor.kill_container(self._running_container) + self.executor.kill_container(entry["container"]) except Exception: pass - self._running_container = None + self._running_clear(task_id) conn = _db_connect(self.db_path) try: cur = conn.execute( @@ -724,8 +799,7 @@ class Worker: proc = self.process_factory(task) else: proc = self.executor.popen_task(task, kind=kind) - self._running_proc = proc - self._running_task_id = task_id + self._running_set(task_id, proc=proc) try: self._set_child_pid(task_id, getattr(proc, "pid", None)) except Exception: @@ -754,8 +828,7 @@ class Worker: pass out, err = proc.communicate() finally: - self._running_proc = None - self._running_task_id = None + self._running_clear(task_id) if cancelled: return {"executor": "local", "task_id": task_id, "status": "cancelled", "stdout": out or "", @@ -839,28 +912,30 @@ class Worker: result = None error = None # Docker 同步路径：dispatch 阻塞在 docker run 里时，取消 - # 线程靠 _running_container 按名杀容器（见 cancel()）。 + # 线程靠容器名杀容器（见 cancel()）。容器名按任务派生， + # 与执行器本身无关，所以这里只记名不碰执行器状态。 if hasattr(self.executor, "container_name"): try: - self._running_container = \ - self.executor.container_name(task) - self._running_task_id = task_id + self._running_set( + task_id, + container=self.executor.container_name(task)) except Exception: - self._running_container = None + pass # Docker 超时取较小值：两个超时都不应静默延长另一个。 - saved_docker_timeout = None + # 注意绝不能改共享的 executor.timeout（并发下两个任务会 + # 互相覆盖对方的超时）；超时只作为本次调用的参数传下去。 + effective_docker_timeout = None is_docker = isinstance(self.executor, DockerExecutor) if is_docker and self.task_timeout is not None: - saved_docker_timeout = self.executor.timeout - if saved_docker_timeout is None: - self.executor.timeout = self.task_timeout + if self.executor.timeout is None: + effective_docker_timeout = self.task_timeout else: try: - self.executor.timeout = min( + effective_docker_timeout = min( float(self.task_timeout), - float(saved_docker_timeout)) + float(self.executor.timeout)) except (TypeError, ValueError): - self.executor.timeout = self.task_timeout + effective_docker_timeout = self.task_timeout try: if self._use_popen(): result = self._execute_via_process( @@ -874,26 +949,36 @@ class Worker: task, kind="run", deadline_seconds=self.task_timeout) else: - result = self.executor.dispatch_task(task) + if is_docker: + result = self.executor.dispatch_task( + task, timeout=effective_docker_timeout) + else: + result = self.executor.dispatch_task(task) combined = (result.get("stdout", "") or "") + \ (result.get("stderr", "") or "") if result["status"] != "completed" and \ _ALREADY_EXISTS_MARKER in combined: - result = self.executor.run_task(task) + if is_docker: + result = self.executor.run_task( + task, timeout=effective_docker_timeout) + else: + result = self.executor.run_task(task) except Exception as exc: # runner 异常也如实写回 failed result = {"executor": "local", "task_id": task_id, "status": "failed", "stdout": "", "returncode": None, "host_status": None, "stderr": ""} error = str(exc) finally: - if is_docker: - self.executor.timeout = saved_docker_timeout # 同步路径执行结束：清掉取消用的容器标记（取消路径里 - # 若已清过，这里再清一次无害）。 - self._running_container = None - if self._running_task_id == task_id and \ - self._running_proc is None: - self._running_task_id = None + # 若已清过，这里再清一次无害）。进程模式的 proc 由 + # _execute_via_process 自己的 finally 清，这里只清 + # 容器项：_running_clear 整个删掉该任务项会连 proc + # 也一起丢，所以这里只在 proc 为空时清容器标记—— + # 实际做法是：如果该任务项里没有 proc，就整项清除； + # 有 proc（进程模式）时不动它，等它自己的 finally。 + entry = self._running_snapshot().get(task_id) or {} + if entry.get("proc") is None: + self._running_clear(task_id) # Docker 超时标记：执行器已在超时时 kill+rm 容器并返回 # failed（stdout 含 timeout），这里把状态统一成 timeout， # 调用方据此置 failed 并记 worker_timeout 审计。 @@ -1185,29 +1270,38 @@ class Worker: fallback_count = 1 self._audit(task["project_id"], "worker_fallback", "%s -> %s" % (self.model, nxt)) - self._set_status(task_id, final) + final_written = self._set_status_if_running(task_id, final) finished_at = time.time() self._write_execution_result( task_id, result_text=stdout, error_text=error, actual_model=actual_model, started_at=started_at, finished_at=finished_at) - self.usage_turns += 1 - # 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构 - # 记 estimated，绝不标 actual；fallback_count 复用同一记录； - # 模型记最后真正跑的那个（fallback 成功时就是 fallback 模型）。 - total = estimate_tokens(task.get("prompt") or "", stdout) - cost_id = self.costs.record( - task_id, task["project_id"], actual_model, - input_tokens=total, output_tokens=0, usd_micros=None, - source="estimated", fallback_count=fallback_count) - self._audit(task["project_id"], - "worker_completed" if final == "completed" - else "worker_failed", task_id) - log_event(self.logger, - "worker_completed" if final == "completed" - else "worker_failed", - task_id=task_id, project_id=task["project_id"], - status=final) + with self._db_lock: + self.usage_turns += 1 + # 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构 + # 记 estimated，绝不标 actual；fallback_count 复用同一记录； + # 模型记最后真正跑的那个（fallback 成功时就是 fallback 模型）。 + # cost 记录共用一条连接，并发下必须持锁写，否则两个线程 + # 共用同一 sqlite 连接会直接报错。 + total = estimate_tokens(task.get("prompt") or "", stdout) + cost_id = self.costs.record( + task_id, task["project_id"], actual_model, + input_tokens=total, output_tokens=0, usd_micros=None, + source="estimated", fallback_count=fallback_count) + if final_written: + self._audit(task["project_id"], + "worker_completed" if final == "completed" + else "worker_failed", task_id) + log_event(self.logger, + "worker_completed" if final == "completed" + else "worker_failed", + task_id=task_id, project_id=task["project_id"], + status=final) + else: + # 已被取消/回收：DB 状态以 interrupted 为准，不再 + # 补一条 worker_completed/failed 审计误导排查。 + final = "interrupted" + self._audit(task["project_id"], "worker_interrupted", task_id) return {"dispatched": True, "task_id": task_id, "host_name": self.executor._host_name(task) if hasattr(self.executor, "_host_name") else task_id, @@ -1221,22 +1315,122 @@ class Worker: 每轮都重新读库，所以上一轮刚 completed 的任务会解锁它的 依赖任务；每轮至少推进一个任务到终态，不会空转死循环。 + + capacity<=1 时走旧串行循环（行为与原来完全一致）。 + capacity>1 时用线程池并行跑：同一 worker 内最多 + capacity 个任务同时执行（capacity 从此是真实并发上限， + 不再只是注册表展示）。认领本身是 CAS，两个线程同时 + 盯上同一任务时只有一个能拿到它，另一个的 run_one 返回 + already_claimed，不计入结果、不死循环。依赖解锁靠 + “等一个在跑任务完成后再认领”：pending 非空但暂无就绪 + 任务时，等一个完成再继续，而不是直接退出。 + 停机信号（_loop_stop/stop_event）只阻止新认领，不会 + 丢弃已提交的任务；排空等待由 run_loop 的 finally 负责。 """ + if self.capacity <= 1: + results = [] + while max_tasks is None or len(results) < max_tasks: + result = self.run_one() + if not result["dispatched"]: + break + results.append(result) + return results + # capacity>1：手动线程池（不用 concurrent.futures）。 + # 原因：本仓库有个单测把 cloud/ 插进 sys.path 后 + # `import queue` 会解析到 cloud/queue.py 而不是标准库， + # concurrent.futures 内部依赖标准库 queue.SimpleQueue， + # 在该单测之后导入它会直接 AttributeError；手动用 + # threading.Condition + 共享计数实现线程池，避开这个坑。 + # 池线程数 = capacity；每个池线程循环认领-执行，认领 + # 本身是 CAS，两个线程盯上同一任务时只有一个能拿到。 + # 依赖解锁：某线程暂无就绪任务时，若仍有线程在跑就等 + # 条件变量（跑完的线程会 notify，可能解锁了依赖）， + # 全部线程都空闲且无就绪任务时才退出。 + # max_tasks 约束：已派发结果数 + 正在跑的任务数达到 + # max_tasks 就不再认领，避免线程池提前提交导致超跑。 + # 停机信号：只阻止新认领，不丢弃已在跑的任务；run_all + # 在停机时立刻返回（后台线程继续跑），run_loop 的 + # finally 做带超时的排空，超时后才杀残留。 + state_lock = threading.Condition() + shared = {"dispatched": 0, "active": 0, "stopped": False} results = [] - while max_tasks is None or len(results) < max_tasks: - result = self.run_one() - if not result["dispatched"]: + results_lock = threading.Lock() + + def _pool_thread(): + while True: + with state_lock: + if shared["stopped"] or self._loop_stop: + return + if max_tasks is not None and \ + shared["dispatched"] + shared["active"] >= max_tasks: + return + try: + task = self.pick_next() + except Exception as exc: + # 认领读库失败（如缺表、锁冲突）不炸池线程： + # 等一轮再试，与 run_loop 的容错一致。 + log_event(self.logger, "worker_pool_pick_error", + message="worker pool pick error: %s" % exc) + state_lock.wait(timeout=0.2) + continue + if task is None: + if shared["active"] > 0: + state_lock.wait(timeout=0.2) + continue + return + shared["active"] += 1 + result = {"dispatched": False} + try: + result = self.run_one(task) + except Exception as exc: + # 单个任务的意外异常绝不能拖垮整个池（串行 + # 模式下 run_loop 也是这么容错的）；该任务已 + # 认领为 running，之后由自动陈旧回收处理。 + log_event(self.logger, "worker_pool_task_error", + message="worker pool task error: %s" % exc) + finally: + with state_lock: + shared["active"] -= 1 + if result["dispatched"]: + shared["dispatched"] += 1 + state_lock.notify_all() + if result["dispatched"]: + with results_lock: + results.append(result) + + threads = [threading.Thread(target=_pool_thread, daemon=True) + for _ in range(self.capacity)] + for t in threads: + t.start() + for t in threads: + # 正常情况等全部池线程结束；停机信号来时不等了， + # 立刻返回给 run_loop 做排空（后台线程是 daemon， + # 且 run_loop finally 会杀残留，不会悬着进程）。 + while t.is_alive(): + t.join(timeout=0.2) + if self._loop_stop: + break + if self._loop_stop: break - results.append(result) return results - def run_loop(self, interval=300, max_seconds=None, stop_event=None): + def run_loop(self, interval=300, max_seconds=None, stop_event=None, + drain_timeout=60, sweep_threshold=300): """常驻循环：每 interval 秒跑一轮 run_all，空闲也不退出。 退出条件：stop_event 被置位、_loop_stop 被信号处理置位、 - 或 max_seconds 到期。退出时若有正在跑的任务（进程模式下 - _running_task_id 非空），先杀进程组、再标 interrupted， - 不让任务永远卡在 running。 + 或 max_seconds 到期。退出时先优雅排空：等正在跑的任务 + 自然结束（最多 drain_timeout 秒），超时后才杀残留进程组/ + 容器并标 interrupted——以前是退出时立刻杀，部署重启会 + 无谓中断一个 16 秒就能跑完的任务。 + + 运行中保障（与并发一起上的两件套）： + - 心跳线程：每 30 秒给所有正在跑的任务刷 last_heartbeat， + 否则长任务（>300 秒）在跑也会被陈旧回收误杀。 + - 自动陈旧回收：每轮结束调一次 sweep_stale_tasks（阈值 + sweep_threshold，默认 300 秒），把死 worker 遗留的 + running 任务标 interrupted 并审计；活任务因心跳新鲜 + 绝不会被扫到。回收失败只记日志，不炸循环。 """ # 常驻启动清理兜底：清掉上一轮残留的 codex- 容器， # 失败只记日志（见 cleanup_stale_containers），不炸循环。 @@ -1249,6 +1443,18 @@ class Worker: self.register() started = time.time() all_results = [] + # 心跳线程：daemon，run_loop 退出即停，不阻塞进程退出。 + hb_stop = threading.Event() + + def _hb_loop(): + while not hb_stop.is_set() and not self._loop_stop: + if stop_event is not None and stop_event.is_set(): + break + self._heartbeat_running_tasks() + hb_stop.wait(30) + + hb_thread = threading.Thread(target=_hb_loop, daemon=True) + hb_thread.start() try: while not self._loop_stop: if stop_event is not None and stop_event.is_set(): @@ -1265,6 +1471,15 @@ class Worker: message="worker_loop_error: %s" % exc) results = [] all_results.extend(results) + # 自动陈旧回收：死 worker 的 running 任务在这里被 + # 标 interrupted；活任务心跳新鲜，不会被扫到。 + if sweep_threshold is not None: + try: + from cloud.server import sweep_stale_tasks + sweep_stale_tasks(self.db_path, sweep_threshold) + except Exception as exc: + log_event(self.logger, "worker_sweep_failed", + message="worker sweep failed: %s" % exc) # 每轮结束写心跳：哪怕本轮一个任务都没派，worker 活着 # 这件事本身也要被 server 看到（文件心跳 + 注册表心跳）。 self._write_heartbeat() @@ -1281,16 +1496,27 @@ class Worker: break time.sleep(0.2) finally: - if self._running_task_id is not None: - task_id = self._running_task_id - if self._running_proc is not None: - self._kill_proc_group(self._running_proc) + hb_stop.set() + # 优雅排空：先等正在跑的任务自然结束，超时后才杀。 + # 串行模式下 run_all 返回时任务已跑完，这里立刻通过； + # 并行模式下停机信号退出时后台线程可能还在跑，先等 + # drain_timeout 秒，超时后按下面杀残留。 + if drain_timeout is not None: + self._wait_for_running(timeout=drain_timeout) + for task_id, entry in self._running_snapshot().items(): + if entry.get("proc") is not None: + self._kill_proc_group(entry["proc"]) try: - self._running_proc.wait(timeout=3) + entry["proc"].wait(timeout=3) except Exception: pass - self._running_proc = None - self._running_task_id = None + if entry.get("container") is not None and \ + hasattr(self.executor, "kill_container"): + try: + self.executor.kill_container(entry["container"]) + except Exception: + pass + self._running_clear(task_id) self._set_status(task_id, "interrupted") return all_results @@ -1360,8 +1586,9 @@ def main(argv=None): help="多 worker 池身份（默认：主机名-进程号；" "生产部署应配稳定 id，否则重启后旧行变 stale）") parser.add_argument("--worker-capacity", type=int, default=1, - help="worker 并发容量声明（注册表里展示用，" - "本轮 run_loop 仍是串行一轮一任务）") + help="worker 并发上限：同一 worker 内最多同时跑" + "几个任务（默认 1，即旧串行行为；>1 时用" + "线程池并行，认领仍是 CAS，成本写库持锁）") parser.add_argument("--docker-image", default="codex-app-server:latest") parser.add_argument("--docker-cpus", default="2") parser.add_argument("--docker-memory", default="2g") @@ -1386,6 +1613,15 @@ def main(argv=None): parser.add_argument("--loop", action="store_true", help="常驻循环，每 --interval 秒跑一轮") parser.add_argument("--interval", type=int, default=300) + parser.add_argument("--drain-timeout", type=float, default=60, + help="停机优雅排空等待上限（秒）：SIGTERM/SIGINT" + "后先等在跑任务自然结束，超时后才杀残留" + "并标 interrupted；默认 60 秒（systemd" + "默认 90 秒停机时限内能完成）") + parser.add_argument("--sweep-threshold", type=float, default=300, + help="自动陈旧回收阈值（秒）：running 任务心跳" + "超过该秒数未更新即被回收；活任务每 30" + "秒刷心跳，不会被误杀") parser.add_argument("--max-seconds", type=float, default=None) args = parser.parse_args(argv) setup_logging() @@ -1426,7 +1662,9 @@ def main(argv=None): try: if args.loop: results = worker.run_loop(interval=args.interval, - max_seconds=args.max_seconds) + max_seconds=args.max_seconds, + drain_timeout=args.drain_timeout, + sweep_threshold=args.sweep_threshold) else: results = worker.run_all(max_tasks=args.max_tasks) finally: