chore: commit all pending distributed changes
This commit is contained in:
Robert Wojciechowski
@@ -1 +0,0 @@
|
||||
|
||||
@@ -74,7 +74,9 @@ def _load_job_routes_module():
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
# torch stub (only needed to satisfy import)
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
@@ -84,7 +86,9 @@ def _load_job_routes_module():
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
# PIL stub (only needed to satisfy import)
|
||||
created_pil_stub = False
|
||||
if "PIL" not in sys.modules:
|
||||
created_pil_stub = True
|
||||
pil_module = types.ModuleType("PIL")
|
||||
image_module = types.ModuleType("PIL.Image")
|
||||
pil_module.Image = image_module
|
||||
@@ -165,6 +169,11 @@ def _load_job_routes_module():
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
if created_pil_stub:
|
||||
sys.modules.pop("PIL.Image", None)
|
||||
sys.modules.pop("PIL", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_media_sync_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
|
||||
package_name = "dist_ms_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
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
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, endpoint="": f"http://localhost{endpoint}"
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in pure-function tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
|
||||
trace_module.trace_debug = lambda *_args, **_kwargs: None
|
||||
trace_module.trace_info = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
class _FormData:
|
||||
def add_field(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.FormData = _FormData
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.media_sync",
|
||||
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)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
ms = _load_media_sync_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# convert_paths_for_platform
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ConvertPathsForPlatformTests(unittest.TestCase):
|
||||
def test_forward_slash_target_normalises_backslashes(self):
|
||||
obj = {"ckpt_name": "C:\\Models\\model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_backslash_target_normalises_forward_slashes(self):
|
||||
obj = {"ckpt_name": "/models/checkpoints/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("\\", result["ckpt_name"])
|
||||
self.assertNotIn("/", result["ckpt_name"])
|
||||
|
||||
def test_relative_media_paths_always_stay_forward_slash(self):
|
||||
"""Relative image/video paths (Comfy annotated style) must not be backslash-ified."""
|
||||
obj = {"image": "subfolder/my_photo.png"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["image"], "subfolder/my_photo.png")
|
||||
|
||||
def test_annotated_relative_media_path_stays_forward_slash(self):
|
||||
obj = {"image": "input/frame.jpg [abc123]"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("/", result["image"])
|
||||
self.assertNotIn("\\", result["image"].split("[")[0])
|
||||
|
||||
def test_non_filename_strings_are_untouched(self):
|
||||
obj = {"prompt": "a beautiful cat", "count": 5}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["prompt"], "a beautiful cat")
|
||||
self.assertEqual(result["count"], 5)
|
||||
|
||||
def test_url_strings_are_untouched(self):
|
||||
obj = {"url": "https://example.com/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["url"], "https://example.com/model.safetensors")
|
||||
|
||||
def test_invalid_separator_returns_obj_unchanged(self):
|
||||
obj = {"ckpt_name": "/models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "|")
|
||||
self.assertEqual(result, obj)
|
||||
|
||||
def test_nested_dict_is_processed_recursively(self):
|
||||
obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_list_items_are_processed_recursively(self):
|
||||
obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"]
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
self.assertEqual(result[1], "plain string")
|
||||
|
||||
def test_non_string_scalar_values_are_untouched(self):
|
||||
obj = {"seed": 42, "enabled": True, "ratio": 1.5}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["seed"], 42)
|
||||
self.assertTrue(result["enabled"])
|
||||
|
||||
def test_absolute_unix_path_to_windows(self):
|
||||
obj = {"lora": "/home/user/loras/my_lora.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertNotIn("/", result["lora"])
|
||||
|
||||
def test_already_normalised_path_is_idempotent(self):
|
||||
obj = {"ckpt": "C:/Models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt"], "C:/Models/model.safetensors")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_media_references
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindMediaReferencesTests(unittest.TestCase):
|
||||
def test_finds_image_input(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.png", refs)
|
||||
|
||||
def test_finds_video_input(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
|
||||
def test_strips_annotation_suffix(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.jpg", refs)
|
||||
self.assertFalse(any("[" in r for r in refs))
|
||||
|
||||
def test_normalises_backslashes_in_path(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("sub/img.png", refs)
|
||||
|
||||
def test_ignores_non_media_text_inputs(self):
|
||||
prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_ignores_node_link_values(self):
|
||||
"""Inputs that are [node_id, slot] lists should be ignored."""
|
||||
prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_deduplicates_same_file_across_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(len(refs), 1)
|
||||
|
||||
def test_returns_sorted_list(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, sorted(refs))
|
||||
|
||||
def test_ignores_non_dict_nodes(self):
|
||||
prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("img.png", refs)
|
||||
|
||||
def test_empty_prompt_returns_empty_list(self):
|
||||
self.assertEqual(ms._find_media_references({}), [])
|
||||
|
||||
def test_multiple_media_types_all_found(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}},
|
||||
"2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("frame.png", refs)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,304 @@
|
||||
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()
|
||||
@@ -143,7 +143,9 @@ def _load_worker_routes_module():
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
@@ -219,6 +221,8 @@ def _load_worker_routes_module():
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
Reference in New Issue
Block a user