Merge feat/runpod-serverless-workers (now including main) into feat/dynamic-image-list-queue
Resolve api/queue_orchestration.py, nodes/v3.py and tests/helpers/v3_runtime_check.py to keep both the worker sink (#99) and the image-list nodes. Stub nodes.image_list in the V3 collector transfer test.
This commit is contained in:
Robert Wojciechowski
@@ -4,6 +4,8 @@ import types
|
||||
import unittest
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
@@ -26,6 +28,17 @@ class _FakeRequest:
|
||||
return self._payload
|
||||
|
||||
|
||||
def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
||||
codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
||||
codec_module = importlib.util.module_from_spec(codec_spec)
|
||||
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
||||
codec_spec.loader.exec_module(codec_module)
|
||||
utils_pkg.transfer_codec = codec_module
|
||||
return codec_module
|
||||
|
||||
|
||||
def _load_job_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
||||
package_name = "dist_api_queue_testpkg"
|
||||
@@ -46,6 +59,7 @@ def _load_job_routes_module():
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
_register_real_transfer_codec(package_name, utils_pkg)
|
||||
|
||||
# aiohttp.web stub
|
||||
created_aiohttp_stub = False
|
||||
@@ -398,5 +412,144 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
job_routes._decode_audio_payload(payload)
|
||||
|
||||
|
||||
class _FakeFileField:
|
||||
def __init__(self, data):
|
||||
self.file = io.BytesIO(data)
|
||||
|
||||
|
||||
class _FakeMultipartRequest:
|
||||
content_type = "multipart/form-data"
|
||||
|
||||
def __init__(self, envelope, parts):
|
||||
from multidict import MultiDict
|
||||
|
||||
fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
|
||||
fields += [("media", _FakeFileField(part)) for part in parts]
|
||||
self._form = MultiDict(fields)
|
||||
|
||||
async def post(self):
|
||||
return self._form
|
||||
|
||||
async def json(self):
|
||||
raise ValueError("not JSON")
|
||||
|
||||
|
||||
class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
codec = job_routes.transfer_codec
|
||||
|
||||
def setUp(self):
|
||||
self.queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
|
||||
|
||||
def _frames(self, count=4, height=20, width=30):
|
||||
rng = np.random.default_rng(1)
|
||||
return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
|
||||
|
||||
def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
|
||||
chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
|
||||
parts = chunk.pop("parts")
|
||||
envelope = {
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"is_last": is_last,
|
||||
**chunk,
|
||||
"start_index": chunk["start_index"] + start_offset,
|
||||
}
|
||||
if audio is not None:
|
||||
envelope["audio"] = audio
|
||||
envelope.update(overrides)
|
||||
return _FakeMultipartRequest(envelope, parts)
|
||||
|
||||
async def _drain(self):
|
||||
items = []
|
||||
while not self.queue.empty():
|
||||
items.append(await self.queue.get())
|
||||
return items
|
||||
|
||||
async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
|
||||
frames = self._frames()
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
|
||||
self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
|
||||
self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
|
||||
for item, frame in zip(items, frames):
|
||||
self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
|
||||
self.assertEqual(item["tensor"].dtype, torch.float32)
|
||||
expected = frame.astype(np.float32) / 255.0
|
||||
self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
|
||||
|
||||
async def test_non_final_chunk_never_marks_worker_done(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "png", is_last=False)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertFalse(any(item["is_last"] for item in await self._drain()))
|
||||
|
||||
async def test_audio_is_attached_only_to_final_frame(self):
|
||||
audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=3), "jpeg", audio=audio)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
|
||||
self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_rejects_frame_count_mismatch(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "ffv1", count=5)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("frames", response.payload["message"])
|
||||
self.assertTrue(self.queue.empty())
|
||||
|
||||
async def test_rejects_bad_envelope(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("format", response.payload["message"])
|
||||
self.assertIn("is_last", response.payload["message"])
|
||||
|
||||
async def test_rejects_invalid_envelope_json(self):
|
||||
response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("envelope", response.payload["message"])
|
||||
|
||||
async def test_unknown_job_returns_404(self):
|
||||
job_routes.prompt_server.distributed_pending_jobs = {}
|
||||
with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
|
||||
|
||||
self.assertEqual(response.status, 404)
|
||||
|
||||
async def test_legacy_json_payload_still_accepted(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 3,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual(len(items), 1)
|
||||
self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -294,3 +294,70 @@ class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fetch_worker_has_node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status, payload):
|
||||
self.status = status
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
if isinstance(self._payload, Exception):
|
||||
raise self._payload
|
||||
return self._payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, response=None, error=None):
|
||||
self.response = response
|
||||
self.error = error
|
||||
self.urls = []
|
||||
|
||||
def get(self, url, timeout=None):
|
||||
self.urls.append(url)
|
||||
if self.error:
|
||||
raise self.error
|
||||
return self.response
|
||||
|
||||
|
||||
class FetchWorkerHasNodeTests(unittest.TestCase):
|
||||
def _probe(self, session):
|
||||
import asyncio
|
||||
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
original = ms.get_client_session
|
||||
ms.get_client_session = _get_session
|
||||
try:
|
||||
return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
|
||||
finally:
|
||||
ms.get_client_session = original
|
||||
|
||||
def test_registered_node_is_reported(self):
|
||||
session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
|
||||
self.assertTrue(self._probe(session))
|
||||
self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
|
||||
|
||||
def test_unknown_node_returns_empty_object(self):
|
||||
# ComfyUI answers 200 {} for classes that are not registered.
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
|
||||
|
||||
def test_http_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
|
||||
|
||||
def test_bad_json_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
|
||||
|
||||
def test_connection_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))
|
||||
|
||||
@@ -0,0 +1,219 @@
|
||||
"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
_TESTS_DIR = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_from(path, name):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _load_modules():
|
||||
real_aiohttp = sys.modules["aiohttp"] # imported inside the isolated test
|
||||
queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
|
||||
collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
|
||||
job_routes = queue_tests._load_job_routes_module()
|
||||
collector = collector_tests._load_collector_module()
|
||||
|
||||
# Error responses must be real HTTP responses here (the unit-test harness
|
||||
# returns plain objects); mirror utils.network.handle_api_error.
|
||||
from aiohttp import web
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
if isinstance(error, list):
|
||||
return web.json_response({"errors": [str(item) for item in error]}, status=status)
|
||||
return web.json_response({"error": str(error)}, status=status)
|
||||
|
||||
job_routes.handle_api_error = handle_api_error
|
||||
# The queue-test harness also fakes aiohttp.web for route modules; the
|
||||
# endpoint must build real responses for the real server used here.
|
||||
job_routes.web = web
|
||||
# The collector harness stubs aiohttp; this test needs the real client.
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
collector.aiohttp = real_aiohttp
|
||||
return job_routes, collector
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_sys_modules():
|
||||
"""The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
|
||||
saved = dict(sys.modules)
|
||||
# Earlier test modules may have left a stub `aiohttp` behind; drop it so the
|
||||
# real package is imported for this test (restored afterwards).
|
||||
for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
|
||||
if getattr(sys.modules[name], "__file__", None) is None:
|
||||
del sys.modules[name]
|
||||
yield
|
||||
sys.modules.clear()
|
||||
sys.modules.update(saved)
|
||||
|
||||
|
||||
async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
|
||||
request_limit, chunk_target, audio=None):
|
||||
"""Send `images` from a real worker node to the real completion route.
|
||||
|
||||
Returns (queued items, sizes of every request body the server received,
|
||||
exception raised by the worker or None).
|
||||
"""
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
body_sizes = []
|
||||
|
||||
@web.middleware
|
||||
async def record_size(request, handler):
|
||||
body_sizes.append(request.content_length)
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
error = None
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
try:
|
||||
await node.send_batch_to_master(
|
||||
images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
except Exception as exc: # surfaced to the caller for assertions
|
||||
error = exc
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items, body_sizes, error
|
||||
|
||||
|
||||
def _noise_images(count, height=360, width=640, seed=97):
|
||||
return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
|
||||
def test_high_entropy_video_stays_under_server_upload_limit(fmt):
|
||||
"""Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(24)
|
||||
limit = 2 * 1024 * 1024
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, fmt,
|
||||
client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
|
||||
|
||||
assert error is None, error
|
||||
assert len(body_sizes) > 1 and max(body_sizes) <= limit
|
||||
assert [item["image_index"] for item in items] == list(range(24))
|
||||
assert [item["is_last"] for item in items] == [False] * 23 + [True]
|
||||
|
||||
|
||||
def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
|
||||
"""A worker configured above the master's limit fails loudly instead of
|
||||
silently resending the whole batch as uncompressed PNG."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(8)
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "ffv1",
|
||||
client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
|
||||
|
||||
assert isinstance(error, RuntimeError)
|
||||
assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
|
||||
assert len(body_sizes) == 1
|
||||
assert items == []
|
||||
|
||||
|
||||
def test_frame_that_cannot_fit_sends_nothing():
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(3)
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "png",
|
||||
client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
|
||||
|
||||
assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
|
||||
assert body_sizes == []
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
|
||||
def test_worker_to_master_over_real_http(fmt):
|
||||
# Import here, not at module level: other test harnesses only install their
|
||||
# aiohttp stubs when the real package has not been imported yet.
|
||||
aiohttp = pytest.importorskip("aiohttp")
|
||||
from aiohttp import web
|
||||
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = torch.rand(5, 24, 36, 3)
|
||||
|
||||
async def scenario():
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000 # several requests for still formats
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
await node.send_batch_to_master(
|
||||
images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items
|
||||
|
||||
items = asyncio.run(scenario())
|
||||
|
||||
assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
|
||||
assert [item["is_last"] for item in items] == [False, False, False, False, True]
|
||||
expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
|
||||
got = torch.cat([item["tensor"] for item in items]).numpy()
|
||||
assert got.shape == expected.shape
|
||||
if fmt in ("png", "ffv1"):
|
||||
assert np.array_equal(got, expected)
|
||||
else:
|
||||
assert np.abs(got - expected).mean() < 0.25 # random noise: only sanity-check lossy codecs
|
||||
Reference in New Issue
Block a user