修复视频子任务空闲节点判断受历史任务影响

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)

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"})

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,

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")