144 lines
7.9 KiB
Python
144 lines
7.9 KiB
Python
import asyncio
|
|
import importlib.util
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PACKAGE = "dist_serverless_runtime_tests"
|
|
for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]:
|
|
module = types.ModuleType(name)
|
|
module.__path__ = [str(path)]
|
|
sys.modules[name] = module
|
|
|
|
def load(name):
|
|
spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py"))
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
runtime = load("serverless_runtime")
|
|
|
|
class Secrets:
|
|
def get(self, worker_id):
|
|
return "private-test-key"
|
|
|
|
class RuntimeTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
|
|
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
|
|
self.manager.poll_seconds = 0
|
|
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
|
|
self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
|
|
|
|
async def asyncTearDown(self):
|
|
await self.manager.cancel()
|
|
self.temp.cleanup()
|
|
|
|
async def test_job_id_persisted_before_status_and_results_injected(self):
|
|
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock())
|
|
async def status(job_id):
|
|
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
|
|
return {"status": "COMPLETED", "output": {"images": []}}
|
|
api.status = status
|
|
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await self.manager.wait_for_execution("execution")
|
|
self.assertEqual(api.submit.await_count, 1)
|
|
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
|
self.assertEqual(item["worker_id"], "pool:1")
|
|
self.assertEqual(item["tensor"], "tensor")
|
|
self.assertTrue(item["is_last"])
|
|
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
|
|
|
async def test_output_error_marks_failed_and_delivers_explicit_error(self):
|
|
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock())
|
|
with patch.object(runtime, "RunpodClient", return_value=api):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await self.manager.wait_for_execution("execution")
|
|
record = self.manager.ledger.all()[0]
|
|
self.assertEqual(record["status"], "FAILED")
|
|
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
|
self.assertIn("serverless_error", item)
|
|
self.assertNotIn("private-test-key", str(record) + str(item))
|
|
|
|
async def test_ambiguous_submit_never_retries(self):
|
|
api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock())
|
|
with patch.object(runtime, "RunpodClient", return_value=api):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await self.manager.wait_for_execution("execution")
|
|
self.assertEqual(api.submit.await_count, 1)
|
|
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
|
|
|
async def test_cancel_targets_owned_job_only(self):
|
|
polled = asyncio.Event()
|
|
async def status(_job):
|
|
polled.set()
|
|
await asyncio.Event().wait()
|
|
api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
|
with patch.object(runtime, "RunpodClient", return_value=api):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await polled.wait()
|
|
await self.manager.cancel(worker_id="other-pool")
|
|
api.cancel.assert_not_awaited()
|
|
await self.manager.cancel(worker_id="pool")
|
|
api.cancel.assert_awaited_once_with("owned-job")
|
|
|
|
async def test_per_pool_concurrency_bounds_submission(self):
|
|
gate = asyncio.Event()
|
|
reached = asyncio.Event()
|
|
count = 0
|
|
async def submit(*args):
|
|
nonlocal count
|
|
count += 1
|
|
reached.set()
|
|
return f"provider-{count}"
|
|
async def status(_job):
|
|
await gate.wait()
|
|
return {"status": "COMPLETED", "output": {"images": []}}
|
|
api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
|
|
second = dict(self.prepared, worker=dict(self.worker, id="pool:2"))
|
|
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
|
await self.manager.start("execution", "prompt", [self.prepared, second])
|
|
await reached.wait()
|
|
await asyncio.sleep(0)
|
|
self.assertEqual(count, 1)
|
|
gate.set()
|
|
await self.manager.wait_for_execution("execution")
|
|
self.assertEqual(count, 2)
|
|
|
|
async def test_cancel_persisted_job_after_master_restart(self):
|
|
self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1)
|
|
api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
|
with patch.object(runtime, "RunpodClient", return_value=api) as client:
|
|
await self.manager.cancel(worker_id="pool")
|
|
api.cancel.assert_awaited_once_with("old-job")
|
|
self.assertEqual(client.call_args.args[1], "old-endpoint")
|
|
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
|
|
|
async def test_master_failure_cancels_job_without_partial_success(self):
|
|
self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}})
|
|
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
|
with patch.object(runtime, "RunpodClient", return_value=api):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await self.manager.wait_for_execution("execution")
|
|
api.submit.assert_not_awaited()
|
|
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
|
|
|
|
async def test_deadline_includes_local_slot_wait(self):
|
|
self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
|
|
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
|
# A blocked slot expires without a billable submission.
|
|
self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0))
|
|
with patch.object(runtime, "RunpodClient", return_value=api):
|
|
await self.manager.start("execution", "prompt", [self.prepared])
|
|
await self.manager.wait_for_execution("execution")
|
|
api.submit.assert_not_awaited()
|
|
self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT")
|