nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: validate complete serverless requests against 10 MB limit

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-10-11 05:52:23 +00:00
parent d9485265ea
commit a281d03c6d
7 changed files with 183 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/orchestration/runpod_client.py
+39 -13
View File
@@ -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:
api/orchestration/serverless_media.py
+19 -5
View File
@@ -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))
api/queue_orchestration.py
+4
View File
@@ -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(
docs/serverless-workers.md
+5 -2
View File
@@ -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.
tests/test_runpod_client.py
+71 -2
View File
@@ -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")))
tests/test_serverless_integration.py
+17
View File
@@ -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, {})
tests/test_serverless_media.py
+28
View File
@@ -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}]