Files
robertvoy-ComfyUI-Distributed/tests/test_serverless_lifecycle.py
T
Robert Wojciechowski 06be0e5dee feat(serverless): replace Concurrent jobs with a single Images per run setting
- Remove the separate Concurrent jobs field; rename Jobs per run to
  Images per run (1-16) with a tooltip explaining that Runpod's max
  workers decides how many GPUs run them.
- Backend uses a fixed per-worker limit of 16 outstanding jobs, so all
  images of a run are submitted at once. Unresolved jobs still count
  against it. The limit is constant, so it never has to change while
  jobs are unresolved.
- Legacy 'concurrency' in saved configs and older browser tabs is
  accepted and ignored/dropped; it no longer triggers the busy guard.
- Revert the enable toggle and show the reason when the backend rejects it.
- Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer
  endpoint-ID error otherwise.
- Confirm before a type change deletes a saved Runpod API key.
- Update docs and tests.
2026-10-13 21:33:02 +00:00

199 lines
11 KiB
Python

"""Billing-safe lifecycle regressions: local completion is not remote completion."""
import asyncio
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
from test_serverless_runtime import runtime, Secrets
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
self.path = Path(self.temp.name) / "jobs.json"
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "timeout_seconds": 30, "type": "serverless"}
self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]}
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
self.client_patch.start()
self.addCleanup(self.client_patch.stop)
self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
self.decode_patch.start()
self.addCleanup(self.decode_patch.stop)
# These regressions exercise capacity holding, so use a limit of one.
self.limit_patch = patch.object(runtime, "PER_WORKER_JOB_LIMIT", 1)
self.limit_patch.start()
self.addCleanup(self.limit_patch.stop)
async def asyncTearDown(self):
await self.manager.cancel()
def second(self, **changes):
return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes))
def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"):
self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id,
status=status, created_at=1, updated_at=1)
async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self):
cancel_started, release_cancel = asyncio.Event(), asyncio.Event()
self.api.submit.side_effect = ["job-1", "job-2"]
self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}]
async def cancel(_job_id):
cancel_started.set()
await release_cancel.wait()
return {"status": "CANCELLED"}
self.api.cancel.side_effect = cancel
await self.manager.start("exec", "prompt", [self.prepared, self.second()])
try:
await asyncio.wait_for(cancel_started.wait(), 1)
await asyncio.sleep(0.01)
self.assertEqual(self.api.submit.await_count, 1)
finally:
release_cancel.set()
await self.manager.wait_for_execution("exec")
self.assertEqual(self.api.submit.await_count, 2)
async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self):
self.api.status.side_effect = runtime.RunpodError("poll failed")
self.api.cancel.return_value = {"status": "IN_PROGRESS"}
await self.manager.start("first", "prompt", [self.prepared])
await self.manager.wait_for_execution("first")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
self.assertTrue(self.manager.pool_is_active("pool"))
await self.manager.start("second", "prompt", [self.second()])
try:
await asyncio.sleep(0.01)
self.assertEqual(self.api.submit.await_count, 1)
self.api.status.side_effect = None
self.api.status.return_value = {"status": "COMPLETED", "output": {}}
await self.manager.reconcile(worker_id="pool")
await asyncio.wait_for(self.manager.wait_for_execution("second"), 1)
self.assertEqual(self.api.submit.await_count, 2)
finally:
self.api.cancel.return_value = {"status": "CANCELLED"}
await self.manager.cancel()
async def test_restarted_manager_reserves_known_unresolved_jobs(self):
self.retained("IN_PROGRESS")
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertFalse(self.manager.pool_is_active("other"))
await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)])
await self.manager.wait_for_execution("exec")
self.api.submit.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT")
async def test_recovery_closes_never_submitted_waiting_records(self):
self.retained("WAITING", job_id=None)
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertFalse(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
self.api.submit.assert_not_awaited()
async def test_recovery_keeps_interrupted_submission_unresolved(self):
self.retained("SUBMITTING", job_id=None)
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
await self.manager.reconcile(worker_id="pool")
self.api.status.assert_not_awaited()
self.api.submit.assert_not_awaited()
async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self):
self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown")
await self.manager.start("exec", "prompt", [self.prepared])
await self.manager.wait_for_execution("exec")
record = self.manager.ledger.all()[0]
self.assertTrue(self.manager.pool_is_active("pool"))
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False)
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True)
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
self.assertFalse(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED")
self.api.submit.assert_awaited_once()
async def test_cancel_already_completed_provider_job_is_terminal(self):
self.retained()
self.api.cancel.return_value = {"status": "COMPLETED"}
await self.manager.cancel(worker_id="pool")
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
self.assertFalse(self.manager.pool_is_active("pool"))
async def test_reconciliation_failure_never_releases_capacity(self):
self.retained()
self.api.status.side_effect = runtime.RunpodError("not reachable")
await self.manager.reconcile(worker_id="pool")
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
async def test_cancel_before_first_task_step_finalizes_local_record(self):
await self.manager.start("exec", "prompt", [self.prepared])
await self.manager.cancel(worker_id="pool")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
self.assertFalse(self.server.distributed_pending_jobs["collector"].empty())
self.assertFalse(self.manager.pool_is_active("pool"))
self.api.submit.assert_not_awaited()
async def test_live_record_cannot_be_acknowledged_or_reconciled(self):
polling = asyncio.Event()
async def status(_job_id):
polling.set()
await asyncio.Event().wait()
self.api.status.side_effect = status
await self.manager.start("exec", "prompt", [self.prepared])
await asyncio.wait_for(polling.wait(), 1)
record = self.manager.ledger.all()[0]
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
await self.manager.reconcile(worker_id="pool")
self.assertEqual(self.api.status.await_count, 1)
self.assertTrue(self.manager.pool_is_active("pool"))
async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self):
polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def status(_job_id):
polling.set()
await asyncio.Event().wait()
async def cancel(_job_id):
cancelling.set()
await release.wait()
return {"status": "CANCELLED"}
self.api.status.side_effect = status
self.api.cancel.side_effect = cancel
await self.manager.start("exec", "prompt", [self.prepared])
await polling.wait()
first = asyncio.create_task(self.manager.cancel(worker_id="pool"))
await cancelling.wait()
second = asyncio.create_task(self.manager.cancel(worker_id="pool"))
await asyncio.sleep(0)
release.set()
await asyncio.gather(first, second)
self.api.cancel.assert_awaited_once_with("job")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
async def test_partial_start_cleanup_does_not_leave_waiting_protection(self):
broken = self.second()
del broken["worker"]["endpoint_id"]
with self.assertRaises(KeyError):
await self.manager.start("exec", "prompt", [self.prepared, broken])
await self.manager.cancel(execution_id="exec")
self.assertFalse(self.manager.pool_is_active("pool"))
self.api.submit.assert_not_awaited()
async def test_public_jobs_never_hides_unresolved_recovery_ids(self):
for index in range(22):
self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index)
for index in range(5):
self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index)
records = self.manager.public_jobs("pool")
self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)})
self.assertTrue(all(record["status"] == "UNKNOWN" for record in records))