From a281d03c6dde77aeed8458fbffb507fff0a5e54b Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 11 Oct 2026 05:52:23 +0000 Subject: [PATCH] fix: validate complete serverless requests against 10 MB limit --- api/orchestration/runpod_client.py | 52 ++++++++++++++----- api/orchestration/serverless_media.py | 24 +++++++-- api/queue_orchestration.py | 4 ++ docs/serverless-workers.md | 7 ++- tests/test_runpod_client.py | 73 ++++++++++++++++++++++++++- tests/test_serverless_integration.py | 17 +++++++ tests/test_serverless_media.py | 28 ++++++++++ 7 files changed, 183 insertions(+), 22 deletions(-) diff --git a/api/orchestration/runpod_client.py b/api/orchestration/runpod_client.py index 64ac958..dbf7c1a 100644 --- a/api/orchestration/runpod_client.py +++ b/api/orchestration/runpod_client.py @@ -6,7 +6,7 @@ import re import aiohttp API_BASE = "https://api.runpod.ai/v2" -MAX_REQUEST_BYTES = 10 * 1024 * 1024 +MAX_REQUEST_BYTES = 10_000_000 _ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z") @@ -18,6 +18,31 @@ class AmbiguousSubmission(RunpodError): """Submission may have been accepted: never retry it automatically.""" +def serialize_request(workflow, images, timeout_seconds): + """Serialize and size-check the exact JSON body used by Runpod ``/run``.""" + payload = { + "input": {"workflow": workflow, "images": images}, + "policy": { + "executionTimeout": timeout_seconds * 1000, + "ttl": timeout_seconds * 1000, + }, + } + body = json.dumps( + payload, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + size = len(body) + if size > MAX_REQUEST_BYTES: + raise ValueError( + f"Serialized serverless request body is {size:,} bytes; Runpod allows at most " + f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw " + "image bytes; reduce image or workflow size." + ) + return body + + def checked_id(value): if not isinstance(value, str) or not _ID.fullmatch(value): raise ValueError("Invalid Runpod endpoint or job ID") @@ -32,19 +57,25 @@ class RunpodClient: raise ValueError("A valid Runpod API key is required") self._api_key = api_key - async def _request(self, method, path, payload=None): + async def _request(self, method, path, payload=None, request_body=None): # Base origin is code-owned, never a user-supplied URL. Do not follow redirects # with credentials, or include response bodies/transport errors in exceptions. url = f"{API_BASE}/{self.endpoint_id}/{path}" attempts = 3 if method == "GET" else 1 for attempt in range(attempts): try: + request_kwargs = { + "headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}, + "timeout": aiohttp.ClientTimeout(total=30), + "allow_redirects": False, + } + if request_body is None: + request_kwargs["json"] = payload + else: + request_kwargs["data"] = request_body async with self.session.request( method, url, - headers={"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}, - json=payload, - timeout=aiohttp.ClientTimeout(total=30), - allow_redirects=False, + **request_kwargs, ) as response: status = response.status if status != 200: @@ -78,13 +109,8 @@ class RunpodClient: return await self._request("GET", "health") async def submit(self, workflow, images, timeout_seconds): - payload = { - "input": {"workflow": workflow, "images": images}, - "policy": {"executionTimeout": timeout_seconds * 1000, "ttl": timeout_seconds * 1000}, - } - if len(json.dumps(payload, allow_nan=False).encode("utf-8")) > MAX_REQUEST_BYTES: - raise ValueError("Serverless request exceeds the 10 MiB safety limit; reduce input media size") - result = await self._request("POST", "run", payload) + body = serialize_request(workflow, images, timeout_seconds) + result = await self._request("POST", "run", request_body=body) try: return checked_id(result.get("id")) except ValueError: diff --git a/api/orchestration/serverless_media.py b/api/orchestration/serverless_media.py index d474eb6..4bf79c1 100644 --- a/api/orchestration/serverless_media.py +++ b/api/orchestration/serverless_media.py @@ -20,7 +20,7 @@ from PIL import Image, UnidentifiedImageError MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024 -MAX_INPUT_TOTAL_BYTES = 7 * 1024 * 1024 +MAX_INPUT_TOTAL_BYTES = 10_000_000 MAX_OUTPUT_FILES = 1024 MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024 MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024 @@ -183,13 +183,23 @@ def _read_limited(path, expected_size): info = os.fstat(handle.fileno()) if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES: raise ServerlessMediaError("Image input exceeds the permitted size") - data = handle.read(MAX_INPUT_IMAGE_BYTES + 1) + if info.st_size > MAX_INPUT_TOTAL_BYTES: + raise ServerlessMediaError( + f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; " + "base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes" + ) + data = handle.read(MAX_INPUT_TOTAL_BYTES + 1) except ServerlessMediaError: raise except OSError: raise ServerlessMediaError("Image input file could not be read safely") from None - if len(data) > MAX_INPUT_IMAGE_BYTES or len(data) != expected_size: - raise ServerlessMediaError("Image input exceeds the permitted size or changed while being read") + if len(data) > MAX_INPUT_TOTAL_BYTES: + raise ServerlessMediaError( + f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; " + "base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes" + ) + if len(data) != expected_size: + raise ServerlessMediaError("Image input changed while being read") return data @@ -229,7 +239,11 @@ async def prepare_images(workflow: dict) -> list[dict]: raise ServerlessMediaError("Image input exceeds the permitted size") total_source_bytes += size if total_source_bytes > MAX_INPUT_TOTAL_BYTES: - raise ServerlessMediaError("Total serverless image inputs exceed the permitted size") + raise ServerlessMediaError( + f"Total raw serverless image input size is {total_source_bytes:,} bytes; " + f"the input read guard is {MAX_INPUT_TOTAL_BYTES:,} bytes. Base64 encoding " + "adds about 33%, and the complete request body is also limited to 10,000,000 bytes" + ) unique_paths[path] = {"size": size, "suffix": suffix} resolved_refs.append((inputs, key, path)) diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index b86e7cb..d66d382 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -165,12 +165,16 @@ async def _prepare_worker_payload( if worker.get("type") == "serverless": from .orchestration.serverless_prompt import compile_serverless_prompt from .orchestration.serverless_media import prepare_images + from .orchestration.runpod_client import serialize_request worker_prompt = apply_participant_overrides( prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids, job_id_map, "", delegate_master, prompt_index, ) workflow, manifest = compile_serverless_prompt(worker_prompt) images = await prepare_images(workflow) + # Fail during preparation, before regular workers or the master are + # dispatched. submit() uses this same serializer for the sent bytes. + serialize_request(workflow, images, worker["timeout_seconds"]) return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images} worker_master_url = build_master_callback_url( diff --git a/docs/serverless-workers.md b/docs/serverless-workers.md index dff76d8..e0b87be 100644 --- a/docs/serverless-workers.md +++ b/docs/serverless-workers.md @@ -109,8 +109,11 @@ with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation. - Downstream post-collector processing remains local; the remote graph contains only the upstream compute required for image saves. - Inputs must be supported images under ComfyUI's input/output/temp roots. -- Image input bytes are bounded to 7 MiB total; the complete request is bounded - to 10 MiB. Large-media/object-storage upload support is not included. +- Raw image input reads are guarded at 10,000,000 bytes total, and the complete + serialized `/run` JSON request body (workflow, base64 images, and policy) must + also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image + budget is lower and depends on workflow size. Large-media/object-storage upload + support is not included. - No remote GPU discovery, endpoint provisioning, model synchronization, automatic scaling changes, or inbound master callback/tunnel is needed. diff --git a/tests/test_runpod_client.py b/tests/test_runpod_client.py index 5b6bc10..fd32239 100644 --- a/tests/test_runpod_client.py +++ b/tests/test_runpod_client.py @@ -1,4 +1,5 @@ import importlib.util +import json import sys import unittest from pathlib import Path @@ -14,6 +15,31 @@ def load_module(): client = load_module() + +def workflow_for_request_size(target_size, timeout_seconds): + """Build an ASCII workflow whose compact UTF-8 request body has exact size.""" + workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}} + + def body_size(candidate): + payload = { + "input": {"workflow": candidate, "images": []}, + "policy": { + "executionTimeout": timeout_seconds * 1000, + "ttl": timeout_seconds * 1000, + }, + } + return len(json.dumps( + payload, allow_nan=False, ensure_ascii=False, separators=(",", ":") + ).encode("utf-8")) + + padding = target_size - body_size(workflow) + if padding < 0: + raise AssertionError("target is smaller than the request fixture") + workflow["1"]["inputs"]["value"] = "x" * padding + if body_size(workflow) != target_size: + raise AssertionError("request fixture did not reach the exact target size") + return workflow + class Response: def __init__(self, status=200, payload=None): self.status, self.payload = status, payload @@ -42,8 +68,10 @@ class ClientTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(result, "job-1") method, url, kwargs = session.calls[0] self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run")) - self.assertEqual(set(kwargs["json"]["input"]), {"workflow", "images"}) - self.assertEqual(kwargs["json"]["policy"]["ttl"], 1800000) + self.assertIsInstance(kwargs["data"], bytes) + sent_payload = json.loads(kwargs["data"]) + self.assertEqual(set(sent_payload["input"]), {"workflow", "images"}) + self.assertEqual(sent_payload["policy"]["ttl"], 1800000) self.assertFalse(kwargs["allow_redirects"]) self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value") @@ -96,3 +124,44 @@ class ClientTests(unittest.IsolatedAsyncioTestCase): with self.assertRaises(ValueError): await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30) self.assertFalse(session.calls) + + async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self): + session = Session([Response(payload={"id": "boundary-job"})]) + api = client.RunpodClient(session, "endpoint123", "key") + timeout_seconds = 37 + workflow = workflow_for_request_size(10_000_000, timeout_seconds) + + result = await api.submit(workflow, [], timeout_seconds) + + self.assertEqual(result, "boundary-job") + sent = session.calls[0][2]["data"] + self.assertIsInstance(sent, bytes) + self.assertEqual(len(sent), 10_000_000) + self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000) + + async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self): + session = Session([Response(payload={"id": "must-not-be-used"})]) + api = client.RunpodClient(session, "endpoint123", "key") + workflow = workflow_for_request_size(10_000_001, 37) + workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:] + + with self.assertRaisesRegex(ValueError, "10,000,001") as ctx: + await api.submit(workflow, [], 37) + + self.assertNotIn("private graph value", str(ctx.exception)) + self.assertFalse(session.calls) + + async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self): + session = Session([Response(payload={"id": "unicode-job"})]) + api = client.RunpodClient(session, "endpoint123", "key") + workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}} + + await api.submit(workflow, [], 41) + + sent = session.calls[0][2]["data"] + expected = json.dumps({ + "input": {"workflow": workflow, "images": []}, + "policy": {"executionTimeout": 41_000, "ttl": 41_000}, + }, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + self.assertEqual(sent, expected) + self.assertGreater(len(sent), len(sent.decode("utf-8"))) diff --git a/tests/test_serverless_integration.py b/tests/test_serverless_integration.py index 79ac7de..b6e17f2 100644 --- a/tests/test_serverless_integration.py +++ b/tests/test_serverless_integration.py @@ -164,3 +164,20 @@ class IntegrationTests(unittest.IsolatedAsyncioTestCase): with self.assertRaisesRegex(ValueError, "invalid local graph"): await self.q.orchestrate_distributed_execution(self.prompt, None, "client") self.manager.start.assert_not_awaited() + + async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self): + media_module = sys.modules[PKG + ".api.orchestration.serverless_media"] + media_module.prepare_images.return_value = [{ + "name": "private-image-name.png", + "image": "x" * 10_000_001, + }] + + with self.assertRaisesRegex(ValueError, "10,000,") as ctx: + await self.q.orchestrate_distributed_execution(self.prompt, None, "client") + + self.assertNotIn("private-image-name.png", str(ctx.exception)) + self.assertNotIn("x" * 100, str(ctx.exception)) + self.q.dispatch_worker_prompt.assert_not_awaited() + self.q.queue_prompt_payload.assert_not_awaited() + self.manager.start.assert_not_awaited() + self.assertEqual(self.ps.distributed_pending_jobs, {}) diff --git a/tests/test_serverless_media.py b/tests/test_serverless_media.py index 054258c..fc88448 100644 --- a/tests/test_serverless_media.py +++ b/tests/test_serverless_media.py @@ -26,6 +26,20 @@ def load_module(): media = load_module() +def load_runpod_module(): + spec = importlib.util.spec_from_file_location( + "dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py" + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +runpod = load_runpod_module() + + def png_bytes(mode="RGB", size=(2, 1), color=None): if color is None: color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30) @@ -201,6 +215,20 @@ class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase): await media.prepare_images(workflow) + async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self): + raw_size = 7 * 1024 * 1024 + 1 + source = self.paths.input / "large-input.png" + source.write_bytes(b"x" * raw_size) + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}} + + packed = await media.prepare_images(workflow) + request_body = runpod.serialize_request(workflow, packed, 30) + + self.assertGreater(raw_size, 7 * 1024 * 1024) + self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size) + self.assertLessEqual(len(request_body), 10_000_000) + + class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase): def manifest(self, prefix="render"): return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}]