198 lines
7.5 KiB
Python
198 lines
7.5 KiB
Python
import asyncio
|
|
import importlib.util
|
|
import sys
|
|
import time
|
|
import types
|
|
import unittest
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
|
|
def _load_job_timeout_module():
|
|
module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
|
|
package_name = "dist_job_timeout_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
|
|
|
|
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
|
|
|
|
config_holder = {"value": {"settings": {}, "workers": []}}
|
|
probe_holder = {"fn": None}
|
|
prompt_server_holder = {
|
|
"value": types.SimpleNamespace(
|
|
distributed_tile_jobs_lock=asyncio.Lock(),
|
|
distributed_pending_tile_jobs={},
|
|
)
|
|
}
|
|
|
|
config_module = types.ModuleType(f"{package_name}.utils.config")
|
|
config_module.load_config = lambda: config_holder["value"]
|
|
sys.modules[f"{package_name}.utils.config"] = config_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.HEARTBEAT_TIMEOUT = 60
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
network_module.build_worker_url = lambda worker: f"http://{worker.get('host', '127.0.0.1')}:{worker.get('port', 8188)}"
|
|
|
|
async def _probe_worker(url, timeout=2.0):
|
|
fn = probe_holder["fn"]
|
|
if fn is None:
|
|
return None
|
|
return await fn(url, timeout)
|
|
|
|
network_module.probe_worker = _probe_worker
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
|
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 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
|
|
|
|
@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
|
|
|
|
job_models_module.BaseJobState = BaseJobState
|
|
job_models_module.ImageJobState = ImageJobState
|
|
job_models_module.TileJobState = TileJobState
|
|
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.job_timeout", 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)
|
|
|
|
module._config_holder = config_holder
|
|
module._probe_holder = probe_holder
|
|
module._prompt_server_holder = prompt_server_holder
|
|
module._ImageJobState = ImageJobState
|
|
module._TileJobState = TileJobState
|
|
return module
|
|
|
|
|
|
jt = _load_job_timeout_module()
|
|
|
|
|
|
class JobTimeoutRequeueTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
jt._prompt_server_holder["value"] = types.SimpleNamespace(
|
|
distributed_tile_jobs_lock=asyncio.Lock(),
|
|
distributed_pending_tile_jobs={},
|
|
)
|
|
jt._config_holder["value"] = {
|
|
"settings": {"worker_timeout_seconds": 5},
|
|
"workers": [{"id": "worker-1", "host": "worker.local", "port": 8188}],
|
|
}
|
|
|
|
async def test_requeues_only_incomplete_dynamic_tasks_for_timed_out_worker(self):
|
|
async def _offline_probe(_url, _timeout):
|
|
return None
|
|
|
|
jt._probe_holder["fn"] = _offline_probe
|
|
prompt_server = jt._prompt_server_holder["value"]
|
|
job_data = jt._ImageJobState("job-1")
|
|
job_data.worker_status["worker-1"] = time.time() - 60.0
|
|
job_data.assigned_to_workers["worker-1"] = [0, 1]
|
|
job_data.completed_images[1] = "done"
|
|
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
|
|
|
requeued = await jt._check_and_requeue_timed_out_workers("job-1", total_tasks=2)
|
|
|
|
self.assertEqual(requeued, 1)
|
|
self.assertEqual(await job_data.pending_images.get(), 0)
|
|
self.assertNotIn("worker-1", job_data.worker_status)
|
|
self.assertEqual(job_data.assigned_to_workers["worker-1"], [])
|
|
|
|
async def test_busy_probe_graces_worker_and_skips_requeue(self):
|
|
async def _busy_probe(_url, _timeout):
|
|
return {"exec_info": {"queue_remaining": 3}}
|
|
|
|
jt._probe_holder["fn"] = _busy_probe
|
|
prompt_server = jt._prompt_server_holder["value"]
|
|
job_data = jt._ImageJobState("job-2")
|
|
old_heartbeat = time.time() - 60.0
|
|
job_data.worker_status["worker-1"] = old_heartbeat
|
|
job_data.assigned_to_workers["worker-1"] = [0]
|
|
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
|
|
|
requeued = await jt._check_and_requeue_timed_out_workers("job-2", total_tasks=1)
|
|
|
|
self.assertEqual(requeued, 0)
|
|
self.assertIn("worker-1", job_data.worker_status)
|
|
self.assertGreaterEqual(job_data.worker_status["worker-1"], old_heartbeat)
|
|
self.assertTrue(job_data.pending_images.empty())
|
|
|
|
async def test_completed_dynamic_task_is_not_requeued(self):
|
|
async def _offline_probe(_url, _timeout):
|
|
return None
|
|
|
|
jt._probe_holder["fn"] = _offline_probe
|
|
prompt_server = jt._prompt_server_holder["value"]
|
|
job_data = jt._ImageJobState("job-3")
|
|
job_data.worker_status["worker-1"] = time.time() - 60.0
|
|
job_data.assigned_to_workers["worker-1"] = [7]
|
|
job_data.completed_images[7] = "complete"
|
|
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
|
|
|
requeued = await jt._check_and_requeue_timed_out_workers("job-3", total_tasks=1)
|
|
|
|
self.assertEqual(requeued, 0)
|
|
self.assertTrue(job_data.pending_images.empty())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|