From 3ca3d31a5efcf1216804b29768610f40dadfd9a2 Mon Sep 17 00:00:00 2001 From: magicwang1111 Date: Fri, 6 Nov 2026 11:24:57 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E8=A7=86=E9=A2=91=E5=AD=90?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1=E7=A9=BA=E9=97=B2=E8=8A=82=E7=82=B9=E5=88=A4?= =?UTF-8?q?=E6=96=AD=E5=8F=97=E5=8E=86=E5=8F=B2=E4=BB=BB=E5=8A=A1=E5=BD=B1?= =?UTF-8?q?=E5=93=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/test_scheduler_store.py | 21 +++++++++++++++++++++ tests/test_task_sessions.py | 11 ++++++++--- utils/scheduler_store.py | 19 +++++++++++++++++++ utils/task_sessions.py | 3 +-- 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/tests/test_scheduler_store.py b/tests/test_scheduler_store.py index fc2793e..67aac9f 100644 --- a/tests/test_scheduler_store.py +++ b/tests/test_scheduler_store.py @@ -40,6 +40,27 @@ class SchedulerStoreTests(unittest.TestCase): self.assertFalse(is_sqlite_busy(sqlite3.OperationalError("disk I/O error"))) self.assertFalse(is_sqlite_busy(RuntimeError("database is locked"))) + def test_idle_workers_ignore_orphan_jobs_and_require_live_compatible_capacity(self): + self.store.register_worker(8180, 100, resource_class="api") + self.store.enqueue(queue_item(0, "orphan"), "old-user", 8180) + self.store.claim(8180, 100, resource_class="api") + # Reproduce a historical running row no longer owned by the live worker. + with closing(self.store._connect()) as connection: + connection.execute("UPDATE scheduler_workers SET active_prompt_id = NULL WHERE port = 8180") + connection.commit() + self.store.register_worker(8181, 101, resource_class="api") + self.store.enqueue(queue_item(1, "active"), "current-user", 8181) + self.store.claim(8181, 101, resource_class="api") + self.store.register_worker(8182, 102, resource_class="gpu") + self.store.register_worker(8183, 103, resource_class="api") + self.store.register_worker(8184, 104, resource_class="default") + with closing(self.store._connect()) as connection: + connection.execute("UPDATE scheduler_workers SET last_heartbeat = 0 WHERE port = 8183") + connection.commit() + self.assertEqual([8180, 8184], self.store.idle_worker_ports("api")) + self.assertEqual("running", self.store.get_job("orphan")["status"]) + self.assertEqual([8182, 8184], self.store.idle_worker_ports("gpu")) + def test_claims_are_unique_across_thirty_workers(self): for index in range(30): self.store.enqueue(queue_item(index, f"job-{index}"), f"user-{index}", 8180) diff --git a/tests/test_task_sessions.py b/tests/test_task_sessions.py index c59b1ca..578bd4f 100644 --- a/tests/test_task_sessions.py +++ b/tests/test_task_sessions.py @@ -29,6 +29,7 @@ class TaskSessionTests(unittest.IsolatedAsyncioTestCase): self.scheduler = Mock( get_job=lambda pid: {"owner_id": pid.split("-")[0], "status": "running"}, worker_ports=lambda: [8180, 8181], + idle_worker_ports=lambda resource: [8181], classify_item=lambda item: "api", resource_concurrency_limits={"api": 2}, max_concurrent_jobs_per_user=2, running_jobs=lambda owner=None, **kw: [{"resource_class": "api", "worker_port": 8180}], @@ -79,9 +80,13 @@ class TaskSessionTests(unittest.IsolatedAsyncioTestCase): {"am_task": "alice-job", "id": "alice"}) async def test_all_workers_busy_rejects_child_instead_of_waiting(self): - self.scheduler.running_jobs = lambda owner=None, **kw: ( - [{"worker_port": 8180}, {"worker_port": 8181}] if kw.get("admin") else [{"resource_class": "api"}] - ) + self.scheduler.idle_worker_ports = lambda resource: [] + request = SimpleNamespace(method="POST", path="/prompt", json=AsyncMock(return_value={"prompt": {}})) + with self.assertRaises(web.HTTPConflict): + await self.sessions.authorize(request, {"am_task": "alice-job", "id": "alice"}) + + async def test_parent_port_cannot_count_as_available_child_worker(self): + self.scheduler.idle_worker_ports = lambda resource: [8183] request = SimpleNamespace(method="POST", path="/prompt", json=AsyncMock(return_value={"prompt": {}})) with self.assertRaises(web.HTTPConflict): await self.sessions.authorize(request, {"am_task": "alice-job", "id": "alice"}) diff --git a/utils/scheduler_store.py b/utils/scheduler_store.py index b7af843..06e953d 100644 --- a/utils/scheduler_store.py +++ b/utils/scheduler_store.py @@ -855,6 +855,25 @@ class SchedulerStore: ).fetchall() return [int(row["port"]) for row in rows] + def idle_worker_ports(self, resource_class: str, stale_seconds: int = 60) -> list[int]: + cutoff = time.time() - max(1, int(stale_seconds)) + with closing(self._connect()) as connection: + rows = connection.execute( + """ + SELECT worker.port + FROM scheduler_workers AS worker + LEFT JOIN scheduler_jobs AS active + ON active.prompt_id = worker.active_prompt_id AND active.status = ? + WHERE worker.last_heartbeat >= ? + AND active.prompt_id IS NULL + AND (worker.resource_class = 'default' OR ? = 'default' + OR worker.resource_class = ?) + ORDER BY worker.port + """, + (RUNNING, cutoff, resource_class, resource_class), + ).fetchall() + return [int(row["port"]) for row in rows] + def upsert_asset( self, asset_id: str, diff --git a/utils/task_sessions.py b/utils/task_sessions.py index 0d8594a..2d85856 100644 --- a/utils/task_sessions.py +++ b/utils/task_sessions.py @@ -70,6 +70,5 @@ class TaskSessions: active = self.scheduler.running_jobs(owner) if limit and sum(job["resource_class"] == resource for job in active) >= limit: raise web.HTTPConflict(reason="Account concurrency is full; parent Agent is waiting for a video child task") - busy_ports = {job["worker_port"] for job in self.scheduler.running_jobs(admin=True)} - if not set(self.scheduler.worker_ports()) - busy_ports: + if not set(self.scheduler.idle_worker_ports(resource)) - {self.server.port}: raise web.HTTPConflict(reason="No free scheduler worker for the video child task")