Robert Wojciechowski
87c0cda9f6
Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.
193 lines
10 KiB
Python
193 lines
10 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", "concurrency": 1, "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)
|
|
|
|
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):
|
|
with self.assertRaises(ValueError):
|
|
await self.manager.start("exec", "prompt", [self.prepared, self.second(concurrency=2)])
|
|
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))
|