Robert Wojciechowski
6b94242110
* feat: add Runpod serverless image collector pools * fix: prevent worker settings clipping and remove preview copy * fix: initialize worker settings before mounting the form * fix: validate complete serverless requests against 10 MB limit * fix: simplify serverless cards and keep status dots circular * fix: match serverless endpoint box to master card * fix: preserve serverless capacity and recovery across job failures Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage. * fix(serverless): disable Clear saved API key while the worker is enabled The checkbox read "Clear saved API key (disable pool first)", but the panel never defines "pool" and saving with it ticked on an enabled worker was rejected by the backend. - Disable the checkbox while the worker is enabled, with a tooltip saying to turn the worker off first. Also disable it when no key is saved. - Use "worker" instead of "pool" in user-facing serverless errors, the serverless guide and the README entry. * feat(serverless): replace Concurrent jobs with a single Images per run setting - Remove the separate Concurrent jobs field; rename Jobs per run to Images per run (1-16) with a tooltip explaining that Runpod's max workers decides how many GPUs run them. - Backend uses a fixed per-worker limit of 16 outstanding jobs, so all images of a run are submitted at once. Unresolved jobs still count against it. The limit is constant, so it never has to change while jobs are unresolved. - Legacy 'concurrency' in saved configs and older browser tabs is accepted and ignored/dropped; it no longer triggers the busy guard. - Revert the enable toggle and show the reason when the backend rejects it. - Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer endpoint-ID error otherwise. - Confirm before a type change deletes a saved Runpod API key. - Update docs and tests.
556 lines
21 KiB
Python
556 lines
21 KiB
Python
import importlib.util
|
|
import sys
|
|
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
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload, status=200):
|
|
self.payload = payload
|
|
self.status = status
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
async def json(self):
|
|
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"
|
|
|
|
# Reset package namespace to avoid stale module state across test runs.
|
|
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
|
|
|
|
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
|
|
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 stub with route decorators
|
|
class _Routes:
|
|
def get(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
def post(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
prompt_server_instance = types.SimpleNamespace(
|
|
routes=_Routes(),
|
|
distributed_jobs_lock=None,
|
|
distributed_pending_jobs={},
|
|
)
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
|
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,
|
|
empty_cache=lambda: None,
|
|
ipc_collect=lambda: None,
|
|
)
|
|
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
|
|
sys.modules["PIL"] = pil_module
|
|
sys.modules["PIL.Image"] = image_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
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
|
image_module.ensure_contiguous = lambda tensor: tensor
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
|
|
|
def _decode_audio_payload(payload):
|
|
if payload is None:
|
|
return None
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Field 'audio' must be an object when provided.")
|
|
|
|
encoded = payload.get("data")
|
|
shape = payload.get("shape")
|
|
dtype = payload.get("dtype", "float32")
|
|
sample_rate = payload.get("sample_rate", 44100)
|
|
if not isinstance(encoded, str) or not encoded.strip():
|
|
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
|
if not isinstance(shape, list) or len(shape) != 3:
|
|
raise ValueError("Field 'audio.shape' must be a 3-item list.")
|
|
if dtype != "float32":
|
|
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
|
try:
|
|
shape_tuple = tuple(int(dim) for dim in shape)
|
|
except Exception as exc:
|
|
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
|
|
|
raw = base64.b64decode(encoded, validate=True)
|
|
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
|
if len(raw) != expected_bytes:
|
|
raise ValueError("Field 'audio.data' byte size mismatch.")
|
|
|
|
waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
|
|
return {"waveform": waveform, "sample_rate": int(sample_rate)}
|
|
|
|
audio_payload_module.decode_audio_payload = _decode_audio_payload
|
|
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_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
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
|
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
|
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
|
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
|
|
|
@dataclass(frozen=True)
|
|
class _QueuePayload:
|
|
prompt: dict
|
|
workflow_meta: object
|
|
client_id: str
|
|
delegate_master: object
|
|
enabled_worker_ids: list
|
|
auto_prepare: bool
|
|
trace_execution_id: object
|
|
|
|
def _parse_queue_request_payload(data):
|
|
if not isinstance(data, dict):
|
|
raise ValueError("Expected a JSON object body")
|
|
prompt = data.get("prompt")
|
|
if not isinstance(prompt, dict):
|
|
raise ValueError("Field 'prompt' must be an object")
|
|
enabled = data.get("enabled_worker_ids")
|
|
if not isinstance(enabled, list):
|
|
raise ValueError("enabled_worker_ids required")
|
|
client_id = data.get("client_id")
|
|
if not isinstance(client_id, str) or not client_id.strip():
|
|
raise ValueError("client_id required")
|
|
return _QueuePayload(
|
|
prompt=prompt,
|
|
workflow_meta=data.get("workflow"),
|
|
client_id=client_id,
|
|
delegate_master=data.get("delegate_master"),
|
|
enabled_worker_ids=enabled,
|
|
auto_prepare=bool(data.get("auto_prepare", True)),
|
|
trace_execution_id=data.get("trace_execution_id"),
|
|
)
|
|
|
|
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
|
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
|
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_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)
|
|
# Keep response assertions independent of real aiohttp import ordering.
|
|
module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
|
|
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
|
|
|
|
|
|
job_routes = _load_job_routes_module()
|
|
|
|
|
|
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
"auto_prepare": True,
|
|
}
|
|
)
|
|
with patch.object(
|
|
job_routes,
|
|
"orchestrate_distributed_execution",
|
|
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
|
):
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
|
self.assertEqual(response.payload.get("number"), 42)
|
|
self.assertEqual(response.payload.get("node_errors"), {})
|
|
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
|
|
|
async def test_distributed_queue_missing_prompt_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("prompt", response.payload.get("message", "").lower())
|
|
|
|
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
|
|
|
|
|
class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
|
def _encoded_audio_payload(self):
|
|
waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
|
|
return {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2, 4],
|
|
"dtype": "float32",
|
|
"data": base64.b64encode(waveform.tobytes()).decode("ascii"),
|
|
}
|
|
|
|
async def test_job_complete_accepts_audio_payload(self):
|
|
queue = asyncio.Queue()
|
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"image": "data:image/png;base64,AAAA",
|
|
"audio": self._encoded_audio_payload(),
|
|
"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)
|
|
queued = await queue.get()
|
|
self.assertEqual(queued["worker_id"], "worker-1")
|
|
self.assertTrue(queued["is_last"])
|
|
self.assertIsNotNone(queued["audio"])
|
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
|
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
|
|
|
async def test_job_complete_accepts_audio_without_image(self):
|
|
queue = asyncio.Queue()
|
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "audio-only-job",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"audio": self._encoded_audio_payload(),
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
decode_image.assert_not_called()
|
|
queued = await queue.get()
|
|
self.assertIsNone(queued["tensor"])
|
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
|
|
|
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
|
|
|
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"image": 123,
|
|
"audio": self._encoded_audio_payload(),
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("image", response.payload.get("message", "").lower())
|
|
|
|
async def test_job_complete_returns_400_for_malformed_audio(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
|
|
|
def test_decode_audio_payload_rejects_bad_shape(self):
|
|
bad = {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2],
|
|
"dtype": "float32",
|
|
"data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
|
|
}
|
|
with self.assertRaises(ValueError):
|
|
job_routes._decode_audio_payload(bad)
|
|
|
|
def test_decode_audio_payload_rejects_bad_dtype(self):
|
|
payload = {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2, 4],
|
|
"dtype": "float16",
|
|
"data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
|
|
}
|
|
with self.assertRaises(ValueError):
|
|
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()
|