305 lines
11 KiB
Python
305 lines
11 KiB
Python
import asyncio
|
|
import importlib.util
|
|
import io
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload, status=200):
|
|
self.payload = payload
|
|
self.status = status
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, json_payload=None, post_payload=None, headers=None, query=None):
|
|
self._json_payload = json_payload
|
|
self._post_payload = post_payload or {}
|
|
self.headers = headers or {}
|
|
self.query = query or {}
|
|
|
|
async def json(self):
|
|
return self._json_payload
|
|
|
|
async def post(self):
|
|
return self._post_payload
|
|
|
|
|
|
class _Routes:
|
|
def post(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
|
|
return _decorator
|
|
|
|
def get(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
|
|
return _decorator
|
|
|
|
|
|
def _load_usdu_routes_module():
|
|
module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
|
|
package_name = "dist_api_usdu_testpkg"
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
api_pkg = types.ModuleType(f"{package_name}.api")
|
|
api_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.api"] = api_pkg
|
|
|
|
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
|
upscale_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
|
|
prompt_server_holder = {
|
|
"value": types.SimpleNamespace(
|
|
distributed_tile_jobs_lock=asyncio.Lock(),
|
|
distributed_pending_tile_jobs={},
|
|
)
|
|
}
|
|
|
|
created_aiohttp_stub = False
|
|
if "aiohttp" not in sys.modules:
|
|
created_aiohttp_stub = True
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
aiohttp_module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
sys.modules["aiohttp"] = aiohttp_module
|
|
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
|
sys.modules["server"] = server_module
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
|
|
async def _handle_api_error(_request, error, status=500):
|
|
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
|
|
|
network_module.handle_api_error = _handle_api_error
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
|
job_store_module.MAX_PAYLOAD_SIZE = 1024
|
|
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
|
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
|
|
|
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
|
|
|
class BaseJobState:
|
|
pass
|
|
|
|
@dataclass
|
|
class TileJobState(BaseJobState):
|
|
multi_job_id: str
|
|
mode: str = field(default="static", init=False)
|
|
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
completed_tasks: dict = field(default_factory=dict)
|
|
worker_status: dict = field(default_factory=dict)
|
|
assigned_to_workers: dict = field(default_factory=dict)
|
|
batch_size: int = 0
|
|
num_tiles_per_image: int = 0
|
|
batched_static: bool = False
|
|
|
|
@dataclass
|
|
class ImageJobState(BaseJobState):
|
|
multi_job_id: str
|
|
mode: str = field(default="dynamic", init=False)
|
|
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
completed_images: dict = field(default_factory=dict)
|
|
worker_status: dict = field(default_factory=dict)
|
|
assigned_to_workers: dict = field(default_factory=dict)
|
|
batch_size: int = 0
|
|
num_tiles_per_image: int = 0
|
|
batched_static: bool = False
|
|
|
|
@property
|
|
def pending_tasks(self):
|
|
return self.pending_images
|
|
|
|
@property
|
|
def completed_tasks(self):
|
|
return self.completed_images
|
|
|
|
job_models_module.BaseJobState = BaseJobState
|
|
job_models_module.TileJobState = TileJobState
|
|
job_models_module.ImageJobState = ImageJobState
|
|
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
|
|
|
parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
|
|
parsers_module._parse_tiles_from_form = lambda _data: []
|
|
sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.api.usdu_routes", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
|
|
if created_aiohttp_stub:
|
|
sys.modules.pop("aiohttp", None)
|
|
|
|
module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
|
|
module._prompt_server_holder = prompt_server_holder
|
|
module._TileJobState = TileJobState
|
|
module._ImageJobState = ImageJobState
|
|
return module
|
|
|
|
|
|
usdu_routes = _load_usdu_routes_module()
|
|
|
|
|
|
class _UploadField:
|
|
def __init__(self, data):
|
|
self.file = io.BytesIO(data)
|
|
|
|
|
|
def _tiny_png_bytes():
|
|
image = Image.new("RGB", (1, 1), (255, 0, 0))
|
|
buf = io.BytesIO()
|
|
image.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
usdu_routes._prompt_server_holder["value"] = types.SimpleNamespace(
|
|
distributed_tile_jobs_lock=asyncio.Lock(),
|
|
distributed_pending_tile_jobs={},
|
|
)
|
|
|
|
async def test_heartbeat_updates_worker_status(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
job_data = usdu_routes._TileJobState("job-1")
|
|
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
|
|
|
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-1"})
|
|
response = await usdu_routes.heartbeat_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("status"), "success")
|
|
self.assertIn("worker-a", job_data.worker_status)
|
|
|
|
async def test_heartbeat_missing_fields_returns_400(self):
|
|
request = _FakeRequest(json_payload={"worker_id": "worker-a"})
|
|
response = await usdu_routes.heartbeat_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("missing", response.payload.get("message", "").lower())
|
|
|
|
async def test_request_image_dynamic_assigns_next_image(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
job_data = usdu_routes._ImageJobState("job-2")
|
|
await job_data.pending_images.put(7)
|
|
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
|
|
|
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-2"})
|
|
response = await usdu_routes.request_image_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("image_idx"), 7)
|
|
self.assertEqual(response.payload.get("estimated_remaining"), 0)
|
|
self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
|
|
self.assertIn("worker-a", job_data.worker_status)
|
|
|
|
async def test_request_image_static_assigns_tile_and_batched_flag(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
job_data = usdu_routes._TileJobState("job-3")
|
|
job_data.batched_static = True
|
|
await job_data.pending_tasks.put(4)
|
|
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
|
|
|
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-3"})
|
|
response = await usdu_routes.request_image_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("tile_idx"), 4)
|
|
self.assertTrue(response.payload.get("batched_static"))
|
|
self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
|
|
|
|
async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
job_data = usdu_routes._TileJobState("job-4")
|
|
prompt_server.distributed_pending_tile_jobs["job-4"] = job_data
|
|
|
|
request = _FakeRequest(
|
|
post_payload={
|
|
"multi_job_id": "job-4",
|
|
"worker_id": "worker-a",
|
|
"batch_size": "0",
|
|
"is_last": "true",
|
|
},
|
|
headers={"content-length": "128"},
|
|
)
|
|
|
|
response = await usdu_routes.submit_tiles_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
queued = await job_data.queue.get()
|
|
self.assertEqual(queued["worker_id"], "worker-a")
|
|
self.assertTrue(queued["is_last"])
|
|
self.assertEqual(queued["tiles"], [])
|
|
|
|
async def test_submit_image_enqueues_processed_image_payload(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
job_data = usdu_routes._ImageJobState("job-5")
|
|
prompt_server.distributed_pending_tile_jobs["job-5"] = job_data
|
|
|
|
request = _FakeRequest(
|
|
post_payload={
|
|
"multi_job_id": "job-5",
|
|
"worker_id": "worker-a",
|
|
"image_idx": "2",
|
|
"full_image": _UploadField(_tiny_png_bytes()),
|
|
"is_last": "false",
|
|
},
|
|
headers={"content-length": "256"},
|
|
)
|
|
|
|
response = await usdu_routes.submit_image_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
queued = await job_data.queue.get()
|
|
self.assertEqual(queued["worker_id"], "worker-a")
|
|
self.assertEqual(queued["image_idx"], 2)
|
|
self.assertIn("image", queued)
|
|
|
|
async def test_job_status_endpoint_reports_ready(self):
|
|
prompt_server = usdu_routes._prompt_server_holder["value"]
|
|
prompt_server.distributed_pending_tile_jobs["job-6"] = usdu_routes._TileJobState("job-6")
|
|
|
|
request = _FakeRequest(query={"multi_job_id": "job-6"})
|
|
response = await usdu_routes.job_status_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertTrue(response.payload.get("ready"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|