From 8ace5731911c99c0560a822b3f3fc91f66e389f6 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:48:16 -0800 Subject: [PATCH] feat!: standardize remote execution on JSON --- README.md | 18 + requirements.txt | 1 - runpod/_serialization.py | 39 +++ runpod/apps/build.py | 1 - runpod/apps/handles.py | 33 +- runpod/apps/images.py | 8 +- runpod/apps/protocol.py | 44 ++- runpod/apps/serialization.py | 63 +--- runpod/apps/targets.py | 40 ++- runpod/apps/tasks.py | 13 +- runpod/endpoint/asyncio/asyncio_runner.py | 3 + runpod/endpoint/runner.py | 2 + runpod/serverless/modules/rp_fastapi.py | 17 +- runpod/serverless/modules/rp_http.py | 17 +- runpod/serverless/modules/rp_job.py | 10 +- tests/test_apps/test_dispatch.py | 28 +- tests/test_apps/test_images.py | 2 - tests/test_apps/test_tasks.py | 22 +- tests/test_serialization.py | 387 ++++++++++++++++++++++ 19 files changed, 620 insertions(+), 128 deletions(-) create mode 100644 runpod/_serialization.py create mode 100644 tests/test_serialization.py diff --git a/README.md b/README.md index 6f5a3fc4d..8e5ad6a9d 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,24 @@ make setup make test ``` +## Remote values + +Queue and task calls use JSON for arguments, return values, and streamed +chunks. Values may contain `None`, booleans, strings, integers, finite +floats, lists, and dictionaries with string keys. Convert tuples to lists, +dates to strings, and model objects to JSON-compatible dictionaries before +calling `.remote()` or `.spawn()`. Store large binary data on a volume or +in object storage and pass its path or URL. + +The same validation applies in live development, deployed workers, and +same-worker remote calls. Unsupported values, circular references, and +non-finite numbers produce explicit errors. `.local()` is a direct Python +call and does not apply the remote serialization contract. Custom ASGI +applications retain their own HTTP response types. + +The SDK and `runpod-sdk-runtime` share the JSON protocol implementation; +use compatible releases of both packages when building worker images. + ## ⚡ | Serverless Worker (SDK) This python package can also be used to create a serverless worker that can be deployed to Runpod as a custom endpoint API. diff --git a/requirements.txt b/requirements.txt index 0bfa95d23..566603c1c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,6 @@ backoff >= 2.2.1 boto3 >= 1.43.118 certifi >= 2026.10.19 click >= 8.5.0 -cloudpickle >= 3.1.2 colorama >= 0.4.6, < 0.4.7 cryptography >= 50.0.2 fastapi[all] >= 0.162.0 diff --git a/runpod/_serialization.py b/runpod/_serialization.py new file mode 100644 index 000000000..3eb58e21d --- /dev/null +++ b/runpod/_serialization.py @@ -0,0 +1,39 @@ +"""strict json values shared by remote callers and workers.""" + +import math +from typing import Any + + +def validate_json(value: Any, label: str = "value") -> None: + """require finite json primitives, lists, and string-keyed dictionaries.""" + active = set() + + def visit(item: Any, path: str) -> None: + if item is None or isinstance(item, (str, bool, int)): + return + if isinstance(item, float): + if not math.isfinite(item): + raise TypeError(f"{path} must be a finite number") + return + if not isinstance(item, (list, dict)): + raise TypeError(f"{path} has unsupported type {type(item).__name__}") + identity = id(item) + if identity in active: + raise TypeError(f"{path} contains a circular reference") + active.add(identity) + try: + if isinstance(item, list): + for index, child in enumerate(item): + visit(child, f"{path}[{index}]") + else: + for key, child in item.items(): + if not isinstance(key, str): + raise TypeError(f"{path} must have string keys") + visit(child, f"{path}[{key!r}]") + finally: + active.remove(identity) + + try: + visit(value, label) + except (TypeError, RecursionError) as exc: + raise TypeError(f"{label} must be json-serializable: {exc}") from exc diff --git a/runpod/apps/build.py b/runpod/apps/build.py index efb91dac6..91575e67f 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -50,7 +50,6 @@ "aiohttp[speedups]", "aiohttp-retry", "backoff", - "cloudpickle", "py-cpuinfo", "requests", "tomli", diff --git a/runpod/apps/handles.py b/runpod/apps/handles.py index b66af1d94..9143504b4 100644 --- a/runpod/apps/handles.py +++ b/runpod/apps/handles.py @@ -31,6 +31,7 @@ def load_model(): ... # worker startup hook Type, ) +from .._serialization import validate_json from .context import Context, block, current_context from .discovery_state import DiscoveryInvocationError, in_discovery from .errors import InvalidResourceError @@ -38,6 +39,7 @@ def load_model(): ... # worker startup hook from .job import Job from .markers import INIT_ATTR, is_init, route_of from .schedule import SCHEDULE_ATTR +from .serialization import require_json_args from .spec import ResourceKind, ResourceSpec, RouteSpec if TYPE_CHECKING: @@ -120,15 +122,34 @@ def local(self, *args: Any, **kwargs: Any) -> Any: def _call_local_sync(self, *args: Any, **kwargs: Any) -> Any: result = self._fn(*args, **kwargs) - return list(result) if inspect.isgenerator(result) else result + if inspect.isgenerator(result): + chunks = [] + for chunk in result: + validate_json(chunk, "return value") + chunks.append(chunk) + result = chunks + if not inspect.isawaitable(result) and not inspect.isasyncgen(result): + validate_json(result, "return value") + return result @staticmethod async def _await_local_result(result: Any) -> Any: if inspect.isawaitable(result): result = await result if inspect.isasyncgen(result): - return [chunk async for chunk in result] - return list(result) if inspect.isgenerator(result) else result + chunks = [] + async for chunk in result: + validate_json(chunk, "return value") + chunks.append(chunk) + result = chunks + elif inspect.isgenerator(result): + chunks = [] + for chunk in result: + validate_json(chunk, "return value") + chunks.append(chunk) + result = chunks + validate_json(result, "return value") + return result def _remote_sync(self, *args: Any, **kwargs: Any) -> Any: if ( @@ -139,6 +160,7 @@ def _remote_sync(self, *args: Any, **kwargs: Any) -> Any: ): self._guard_discovery() self.spec.validate() + require_json_args(self._fn, args, kwargs) result = self._call_local_sync(*args, **kwargs) if inspect.isawaitable(result) or inspect.isasyncgen(result): return block(self._await_local_result(result)) @@ -148,6 +170,7 @@ def _remote_sync(self, *args: Any, **kwargs: Any) -> Any: async def _remote_async(self, *args: Any, **kwargs: Any) -> Any: self._guard_discovery() self.spec.validate() + require_json_args(self._fn, args, kwargs) ctx = current_context() if ctx is Context.WORKER and self._is_current_worker(): @@ -169,15 +192,18 @@ async def _stream_async(self, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: self._guard_discovery() self.spec.validate() self._guard_generator() + require_json_args(self._fn, args, kwargs) ctx = current_context() if ctx is Context.WORKER and self._is_current_worker(): gen = self._fn(*args, **kwargs) if inspect.isasyncgen(gen): async for chunk in gen: + validate_json(chunk, "stream chunk") yield chunk else: for chunk in gen: + validate_json(chunk, "stream chunk") yield chunk return @@ -208,6 +234,7 @@ def _apply_options(self, payload: Dict[str, Any]) -> Dict[str, Any]: async def _spawn_async(self, *args: Any, **kwargs: Any) -> Any: self._guard_discovery() + require_json_args(self._fn, args, kwargs) target = await self._app._resolve(self.spec) payload = target.build_payload(self._fn, self.spec, args, kwargs) data = await target.submit(self._apply_options(payload)) diff --git a/runpod/apps/images.py b/runpod/apps/images.py index fb00e1947..343f58fa7 100644 --- a/runpod/apps/images.py +++ b/runpod/apps/images.py @@ -38,11 +38,9 @@ def runtime_tag() -> str: def local_python_version() -> str: """the local interpreter's version, which the worker must match. - dev sessions and tasks ship cloudpickle payloads produced on the - client interpreter; those are not reliably portable across python - versions, so a silent fallback would surface as cryptic - deserialization failures on the worker. mismatches fail here, - loudly, instead. + live execution ships python source from the client. matching the + interpreter keeps syntax and language behavior consistent between + local development and the worker. """ version = f"{sys.version_info.major}.{sys.version_info.minor}" if version not in SUPPORTED_PYTHON_VERSIONS: diff --git a/runpod/apps/protocol.py b/runpod/apps/protocol.py index bc9f34bbf..255afae97 100644 --- a/runpod/apps/protocol.py +++ b/runpod/apps/protocol.py @@ -1,21 +1,20 @@ """wire protocol for remote function execution. this is the contract between the sdk and the worker runtime images. -the worker repo imports these shapes from the runpod package so client -and worker can never drift. +the worker runtime and client share these shapes and validation rules. live mode ships function source with every request; deployed mode omits the code (the worker unpacked the build and resolves the function by -name). args/kwargs are base64 cloudpickle unless serialization_format -is "json". +name). arguments, results, and stream chunks are json values. """ -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, fields from typing import Any, Dict, List, Optional +from .._serialization import validate_json + EXECUTION_FUNCTION = "function" -FORMAT_CLOUDPICKLE = "cloudpickle" FORMAT_JSON = "json" @@ -31,10 +30,28 @@ class FunctionRequest: system_dependencies: Optional[List[str]] = None execution_type: str = EXECUTION_FUNCTION accelerate_downloads: bool = True - serialization_format: str = FORMAT_CLOUDPICKLE + serialization_format: str = FORMAT_JSON + + def validate(self) -> None: + if self.serialization_format != FORMAT_JSON: + raise ValueError("serialization_format must be 'json'") + if not isinstance(self.args, list) or not isinstance(self.kwargs, dict): + raise TypeError("args must be a list and kwargs must be a dictionary") + validate_json( + {f.name: getattr(self, f.name) for f in fields(self)}, "function request" + ) + + @classmethod + def from_input(cls, data: Dict[str, Any]) -> "FunctionRequest": + if not isinstance(data, dict): + raise TypeError("function request must be a dictionary") + request = cls(**{f.name: data[f.name] for f in fields(cls) if f.name in data}) + request.validate() + return request def to_input(self) -> Dict[str, Any]: """serialize as the job input dict.""" + self.validate() data = asdict(self) return {k: v for k, v in data.items() if v is not None} @@ -44,16 +61,23 @@ class FunctionResponse: """the worker's result for one function execution.""" success: bool - result: Optional[str] = None json_result: Any = None error: Optional[str] = None stdout: Optional[str] = None @classmethod def from_output(cls, output: Dict[str, Any]) -> "FunctionResponse": + if not isinstance(output, dict) or not isinstance(output.get("success"), bool): + raise TypeError("function response must contain a boolean success field") + if output.get("serialization_format", FORMAT_JSON) != FORMAT_JSON: + raise ValueError("serialization_format must be 'json'") + if output.get("result") is not None: + raise ValueError("function responses must use json_result") + if output["success"] and "json_result" not in output: + raise ValueError("successful function response is missing json_result") + validate_json(output, "function response") return cls( - success=output.get("success", False), - result=output.get("result"), + success=output["success"], json_result=output.get("json_result"), error=output.get("error"), stdout=output.get("stdout"), diff --git a/runpod/apps/serialization.py b/runpod/apps/serialization.py index 425189e1a..1ced658cc 100644 --- a/runpod/apps/serialization.py +++ b/runpod/apps/serialization.py @@ -1,63 +1,32 @@ -"""argument and result serialization for remote execution. - -args/kwargs cross the wire as base64-encoded cloudpickle strings, the -format the worker runtime images expect. function source is extracted -with decorators stripped so the worker can exec it standalone. -""" +"""json argument validation, response decoding, and function source shipping.""" import ast -import base64 import inspect import os import textwrap -from typing import Any, Callable, Dict, List - -import cloudpickle - - -def serialize_arg(arg: Any) -> str: - return base64.b64encode(cloudpickle.dumps(arg)).decode("utf-8") - - -def serialize_args(args: tuple) -> List[str]: - return [serialize_arg(a) for a in args] - +from typing import Any, Callable -def serialize_kwargs(kwargs: dict) -> Dict[str, str]: - return {k: serialize_arg(v) for k, v in kwargs.items()} +from .._serialization import validate_json +from .errors import RemoteExecutionError +from .protocol import FunctionResponse -def deserialize_result(result_b64: str) -> Any: +def unwrap_function_response(output: dict) -> Any: try: - return cloudpickle.loads(base64.b64decode(result_b64)) - except ModuleNotFoundError as exc: - from .errors import RemoteExecutionError - + response = FunctionResponse.from_output(output) + except (TypeError, ValueError) as exc: + raise RemoteExecutionError(f"invalid remote response: {exc}") from exc + if not response.success: raise RemoteExecutionError( - f"the remote result contains an object from the " - f"'{exc.name}' package, which is not installed locally. " - f"return plain python types (e.g. str(...)) or install " - f"'{exc.name}' on this machine." - ) from exc + f"remote execution failed: {response.error or 'unknown'}" + ) + return response.json_result def require_json_args(fn: Callable, args: tuple, kwargs: dict) -> None: - """reject arguments that cannot cross the queue data plane. - - queue jobs are json on the wire in every mode, so non-json - arguments fail at call time with a pointed error instead of a - worker-side decode failure. - """ - import json - - try: - json.dumps(list(args)) - json.dumps(kwargs) - except (TypeError, ValueError) as exc: - raise TypeError( - f"{fn.__name__}() arguments must be json-serializable " - f"(queue jobs are plain json): {exc}" - ) from exc + """validate queue and task arguments before remote execution.""" + validate_json(list(args), f"{fn.__name__}() arguments") + validate_json(kwargs, f"{fn.__name__}() keyword arguments") def get_function_source(fn: Callable) -> str: diff --git a/runpod/apps/targets.py b/runpod/apps/targets.py index 28e720513..02e6d8d3b 100644 --- a/runpod/apps/targets.py +++ b/runpod/apps/targets.py @@ -13,9 +13,10 @@ import aiohttp +from .._serialization import validate_json from ..user_agent import USER_AGENT from .errors import EndpointNotFound, RemoteExecutionError -from .protocol import FORMAT_JSON, FunctionRequest, FunctionResponse +from .protocol import FORMAT_JSON, FunctionRequest from .spec import ResourceSpec from .utils.events import emit from .utils.network import api_key as _api_key @@ -116,6 +117,7 @@ def merge_job_options(base: Dict[str, Any], new: Dict[str, Any]) -> Dict[str, An def unwrap_job_output(data: Dict[str, Any]) -> Any: """extract output from a queue response, raising on failure.""" + validate_json(data, "job response") output = data.get("output", data) if data.get("status") == "FAILED" or data.get("error"): err = data.get("error") @@ -148,6 +150,9 @@ class InvocationTarget(ABC): def build_payload( self, fn: Callable, spec: ResourceSpec, args: tuple, kwargs: dict ) -> Dict[str, Any]: + from .serialization import require_json_args + + require_json_args(fn, args, kwargs) return {"input": args_to_input(fn, args, kwargs)} def unwrap(self, data: Dict[str, Any]) -> Any: @@ -218,6 +223,7 @@ async def _request_json( retry_not_found: bool = False, ) -> Dict[str, Any]: """http json call with exponential backoff on transient failures.""" + validate_json(payload, "request body") retry_transient = method.upper() in {"GET", "HEAD", "OPTIONS"} client_timeout = aiohttp.ClientTimeout(total=timeout) @@ -452,7 +458,9 @@ async def stream(self, job_id: str, *, timeout: float) -> AsyncIterator[Any]: if data.get("status") == "FAILED" or data.get("error"): unwrap_job_output(data) for chunk in data.get("stream") or []: - yield chunk.get("output") + output = chunk.get("output") + validate_json(output, "stream chunk") + yield output if data.get("status") in FINAL_STATUSES: return @@ -482,6 +490,7 @@ async def http( timeout: float = DEFAULT_TIMEOUT_SECONDS, ) -> Any: """plain http call through the endpoint's lb subdomain.""" + validate_json(body, "request body") url = f"https://{self.endpoint_id}.{_lb_domain()}{path}" client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: @@ -642,9 +651,6 @@ def build_payload( ) -> Dict[str, Any]: from .serialization import get_function_source, require_json_args - # queue arguments are plain json, the same contract deployed - # workers serve; non-json arguments fail here in dev exactly as - # they would against a deployed endpoint require_json_args(fn, args, kwargs) request = FunctionRequest( function_name=fn.__name__, @@ -675,16 +681,10 @@ def unwrap(self, data: Dict[str, Any]) -> Any: def _unwrap_response(self, output: Any) -> Any: if isinstance(output, dict) and "success" in output: - response = FunctionResponse.from_output(output) - if not response.success: - raise RemoteExecutionError( - f"remote execution failed: {response.error or 'unknown'}" - ) - if response.result is not None: - from .serialization import deserialize_result + from .serialization import unwrap_function_response - return deserialize_result(response.result) - return response.json_result + return unwrap_function_response(output) + validate_json(output, "remote result") return output def _monitor(self): @@ -822,20 +822,18 @@ def __init__( def build_payload( self, fn: Callable, spec: ResourceSpec, args: tuple, kwargs: dict ) -> Dict[str, Any]: - from .serialization import ( - get_function_source, - serialize_args, - serialize_kwargs, - ) + from .serialization import get_function_source, require_json_args + require_json_args(fn, args, kwargs) request = FunctionRequest( function_name=fn.__name__, function_code=get_function_source(fn), - args=serialize_args(args), - kwargs=serialize_kwargs(kwargs), + args=list(args), + kwargs=dict(kwargs), dependencies=spec.dependencies, system_dependencies=spec.system_dependencies, accelerate_downloads=spec.accelerate_downloads, + serialization_format=FORMAT_JSON, ) return request.to_input() diff --git a/runpod/apps/tasks.py b/runpod/apps/tasks.py index e6a6419b4..3691e80b2 100644 --- a/runpod/apps/tasks.py +++ b/runpod/apps/tasks.py @@ -19,9 +19,9 @@ import aiohttp +from .._serialization import validate_json from .api import AppsApiClient, is_capacity_error from .context import CLEANUP_TIMEOUT -from .errors import RemoteExecutionError from .spec import ResourceSpec log = logging.getLogger(__name__) @@ -319,6 +319,7 @@ async def _post(self, path: str, request: Dict[str, Any]) -> None: or timeout < 0 ): raise ValueError("task timeout must be a finite non-negative number") + validate_json(request, "request body") url = f"{_proxy_url(self.pod_id)}{path}" attempts = 6 async with aiohttp.ClientSession() as session: @@ -423,15 +424,9 @@ async def terminate(self) -> None: def unwrap_task_response(response: Dict[str, Any]) -> Any: """extract the function result from a runner response.""" - if not response.get("success"): - raise RemoteExecutionError( - f"task execution failed: {response.get('error', 'unknown')}" - ) - if response.get("result") is not None: - from .serialization import deserialize_result + from .serialization import unwrap_function_response - return deserialize_result(response["result"]) - return response.get("json_result") + return unwrap_function_response(response) class TaskJob: diff --git a/runpod/endpoint/asyncio/asyncio_runner.py b/runpod/endpoint/asyncio/asyncio_runner.py index cb75144d1..17f0aae73 100644 --- a/runpod/endpoint/asyncio/asyncio_runner.py +++ b/runpod/endpoint/asyncio/asyncio_runner.py @@ -157,6 +157,9 @@ async def run(self, endpoint_input: dict) -> Job: Returns: Newly created job """ + from runpod._serialization import validate_json + + validate_json(endpoint_input, "job input") async with self.session.post( self.endpoint_url, headers=self.headers, json={"input": endpoint_input} ) as resp: diff --git a/runpod/endpoint/runner.py b/runpod/endpoint/runner.py index d8c487599..3e736c790 100644 --- a/runpod/endpoint/runner.py +++ b/runpod/endpoint/runner.py @@ -8,6 +8,7 @@ import requests from requests.adapters import HTTPAdapter, Retry +from runpod._serialization import validate_json from runpod.endpoint.helpers import ( API_KEY_NOT_SET_MSG, FINAL_STATES, @@ -75,6 +76,7 @@ def _request( requests.HTTPError: If the response contains an unsuccessful status code. """ url = f"{self.endpoint_url_base}/{endpoint}" + validate_json(data, "request body") response = self.rp_session.request( method, url, headers=self.headers, json=data, timeout=timeout ) diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index f5b7c0c6e..490702163 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -267,7 +267,6 @@ def __init__(self, config: dict[str, Any]): "/runsync", self._sim_runsync, methods=["POST"], - response_model_exclude_none=True, summary="Mimics the behavior of the runsync endpoint.", description=RUNSYNC_DESCRIPTION, tags=["Synchronously Submit Request & Get Job Results"], @@ -285,7 +284,6 @@ def __init__(self, config: dict[str, Any]): "/status/{job_id}", self._sim_status, methods=["POST"], - response_model_exclude_none=True, summary="Mimics the behavior of the status endpoint.", description=STATUS_DESCRIPTION, tags=["Check Job Results"], @@ -348,6 +346,9 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput: generator_output = run_job_generator(self.config["handler"], job.__dict__) job_output = {"output": []} async for stream_output in generator_output: + if stream_output.get("error"): + job_output = stream_output + break job_output["output"].append(stream_output["output"]) else: job_output = await run_job(self.config["handler"], job.__dict__) @@ -366,7 +367,7 @@ async def _sim_runsync(self, job_request: DefaultRequest) -> JobOutput: thread.start() return jsonable_encoder( - {"id": job.id, "status": "COMPLETED", "output": job_output["output"]} + {"id": job.id, "status": "COMPLETED", "output": job_output.get("output", {})} ) # ---------------------------------- stream ---------------------------------- # @@ -384,6 +385,11 @@ async def _sim_stream(self, job_id: str) -> StreamOutput: generator_output = run_job_generator(self.config["handler"], job.__dict__) stream_accumulator = [] async for stream_output in generator_output: + if stream_output.get("error"): + job_list.remove(job.id) + return jsonable_encoder( + {"id": job_id, "status": "FAILED", "error": stream_output["error"]} + ) stream_accumulator.append({"output": stream_output["output"]}) else: return jsonable_encoder( @@ -423,6 +429,9 @@ async def _sim_status(self, job_id: str) -> JobOutput: generator_output = run_job_generator(self.config["handler"], job.__dict__) job_output = {"output": []} async for stream_output in generator_output: + if stream_output.get("error"): + job_output = stream_output + break job_output["output"].append(stream_output["output"]) else: job_output = await run_job(self.config["handler"], job.__dict__) @@ -443,5 +452,5 @@ async def _sim_status(self, job_id: str) -> JobOutput: thread.start() return jsonable_encoder( - {"id": job_id, "status": "COMPLETED", "output": job_output["output"]} + {"id": job_id, "status": "COMPLETED", "output": job_output.get("output", {})} ) diff --git a/runpod/serverless/modules/rp_http.py b/runpod/serverless/modules/rp_http.py index 3d82d35b6..1f9dfeed7 100644 --- a/runpod/serverless/modules/rp_http.py +++ b/runpod/serverless/modules/rp_http.py @@ -8,6 +8,7 @@ from aiohttp import ClientError from aiohttp_retry import FibonacciRetry, RetryClient +from runpod._serialization import validate_json from runpod.http_client import ClientSession from runpod.serverless.modules.rp_logger import RunPodLogger @@ -57,7 +58,19 @@ async def _handle_result( try: session.headers["X-Request-ID"] = job["id"] - serialized_job_data = json.dumps(job_data, ensure_ascii=False) + try: + validate_json(job_data, "job result") + serialized_job_data = json.dumps( + job_data, ensure_ascii=False, allow_nan=False + ) + except (TypeError, ValueError, RecursionError) as err: + error = { + "error_type": "SerializationError", + "error_message": str(err), + } + serialized_job_data = json.dumps({"error": json.dumps(error)}) + url_template = JOB_DONE_URL + log.error(f"Error while returning job result. | {err}", job["id"]) is_stream = "true" if is_stream else "false" url = url_template.replace("$ID", job["id"]) + f"&isStream={is_stream}" @@ -68,7 +81,7 @@ async def _handle_result( except ClientError as err: log.error(f"Failed to return job results. | {err}", job["id"]) - except (TypeError, RuntimeError) as err: + except (TypeError, ValueError, RuntimeError) as err: log.error(f"Error while returning job result. | {err}", job["id"]) finally: diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index 108efc4a6..fc1d40c0f 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/runpod/serverless/modules/rp_job.py @@ -10,6 +10,7 @@ import aiohttp +from runpod._serialization import validate_json from runpod.http_client import ClientSession, TooManyRequests from runpod.serverless.modules.rp_logger import RunPodLogger @@ -256,12 +257,14 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]: with capture(enabled=handler_capture_enabled()) as cap: try: + validate_json(job.get("input"), "job input") handler_return = handler(job) job_output = ( await handler_return if inspect.isawaitable(handler_return) else handler_return ) + validate_json(job_output, "return value") log.debug(f"Handler output: {job_output}", job["id"]) @@ -281,7 +284,9 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]: else: run_result = {"output": job_output} - if run_result.get("output") == {}: + if run_result.get("output") == {} and ( + "error" in run_result or "stopPod" in run_result + ): run_result.pop("output") check_return_size(run_result) # Checks the size of the return body. @@ -324,14 +329,17 @@ async def run_job_generator( with capture(enabled=handler_capture_enabled()) as cap: try: + validate_json(job.get("input"), "job input") job_output = handler(job) if is_async_gen: async for output_partial in job_output: + validate_json(output_partial, "stream chunk") log.debug(f"Async Generator output: {output_partial}", job["id"]) yield {"output": output_partial} else: for output_partial in job_output: + validate_json(output_partial, "stream chunk") log.debug(f"Generator output: {output_partial}", job["id"]) yield {"output": output_partial} diff --git a/tests/test_apps/test_dispatch.py b/tests/test_apps/test_dispatch.py index 6a2543a83..d9ed409c1 100644 --- a/tests/test_apps/test_dispatch.py +++ b/tests/test_apps/test_dispatch.py @@ -546,19 +546,29 @@ def test_whole_module_ships(self, tmp_path): finally: sys.modules.pop("shipmod", None) - def test_missing_module_on_deserialize_is_actionable(self): + def test_pickled_result_is_rejected_without_importing_remote_modules( + self, monkeypatch + ): import base64 + import builtins + import pickle - import cloudpickle - import pytest - - from runpod.apps.errors import RemoteExecutionError - from runpod.apps.serialization import deserialize_result + from runpod.apps.serialization import unwrap_function_response class FakePickle: def __reduce__(self): return (__import__, ("nonexistent_pkg_xyz",)) - payload = base64.b64encode(cloudpickle.dumps(FakePickle())).decode() - with pytest.raises(RemoteExecutionError, match="nonexistent_pkg_xyz"): - deserialize_result(payload) + payload = base64.b64encode(pickle.dumps(FakePickle())).decode() + real_import = builtins.__import__ + imported = [] + + def track_import(name, *args, **kwargs): + if name == "nonexistent_pkg_xyz": + imported.append(name) + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", track_import) + with pytest.raises(RemoteExecutionError, match="must use json_result"): + unwrap_function_response({"success": True, "result": payload}) + assert imported == [] diff --git a/tests/test_apps/test_images.py b/tests/test_apps/test_images.py index d9c8a88bc..915353627 100644 --- a/tests/test_apps/test_images.py +++ b/tests/test_apps/test_images.py @@ -63,8 +63,6 @@ def test_returns_supported_version(self): assert local_python_version() in SUPPORTED_PYTHON_VERSIONS def test_unsupported_local_python_fails_loudly(self, monkeypatch): - # cloudpickle payloads are version-bound; a silent fallback - # would break at deserialization time on the worker class FakeVersion: major, minor = 3, 9 diff --git a/tests/test_apps/test_tasks.py b/tests/test_apps/test_tasks.py index 78b761152..04853971d 100644 --- a/tests/test_apps/test_tasks.py +++ b/tests/test_apps/test_tasks.py @@ -3,9 +3,9 @@ import asyncio import base64 import json +import pickle from unittest.mock import AsyncMock, MagicMock, patch -import cloudpickle import pytest from runpod.apps import App @@ -31,14 +31,6 @@ def clean_registry(): _clear_registry() -def _b64(value): - return base64.b64encode(cloudpickle.dumps(value)).decode("utf-8") - - -def _unb64(value): - return cloudpickle.loads(base64.b64decode(value)) - - class TestPodInput: def _spec(self, **kw): defaults = dict(kind=ResourceKind.TASK, name="t") @@ -113,8 +105,10 @@ def test_cpu5_uses_supported_datacenters(self): class TestUnwrapTaskResponse: - def test_success_cloudpickle(self): - assert unwrap_task_response({"success": True, "result": _b64(42)}) == 42 + def test_pickled_result_is_rejected(self): + payload = base64.b64encode(pickle.dumps(42)).decode("utf-8") + with pytest.raises(RemoteExecutionError, match="must use json_result"): + unwrap_task_response({"success": True, "result": payload}) def test_success_json(self): assert unwrap_task_response({"success": True, "json_result": {"a": 1}}) == { @@ -138,7 +132,9 @@ def t(x): payload = target.build_payload(t._fn, t.spec, (5,), {}) assert payload["function_name"] == "t" assert "def t(x):" in payload["function_code"] - assert _unb64(payload["args"][0]) == 5 + assert payload["args"] == [5] + assert payload["kwargs"] == {} + assert payload["serialization_format"] == "json" assert payload["dependencies"] == ["numpy"] def test_task_remote_runs_past_one_hour_and_cleans_up(self): @@ -159,7 +155,7 @@ async def poll_result(): if clock[0] == 0: clock[0] = 7200.0 return None - return {"success": True, "result": _b64(10)} + return {"success": True, "json_result": 10} async def delete(pod_id): pods.remove(pod_id) diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 000000000..662ee5a50 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,387 @@ +"""shared json contracts for queues, tasks, and serverless workers.""" + +import json +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from runpod._serialization import validate_json +from runpod.apps import App +from runpod.apps.app import _clear_registry +from runpod.apps.errors import RemoteExecutionError +from runpod.apps.protocol import FunctionRequest, FunctionResponse +from runpod.apps.spec import ResourceKind, ResourceSpec +from runpod.apps.targets import LiveTarget, PodTarget, SentinelTarget +from runpod.apps.tasks import unwrap_task_response +from runpod.serverless.modules import rp_fastapi, rp_http, rp_job + + +def identity(value): + return value + + +@pytest.mark.parametrize( + "value", [None, True, False, 0, -1, 1.25, "hello", [], {}, {"x": [None, 2]}] +) +def test_json_values_roundtrip_without_coercion(value): + validate_json(value) + assert json.loads(json.dumps(value, allow_nan=False)) == value + + +@pytest.mark.parametrize( + "value", + [ + b"bytes", + (1, 2), + {1: "key"}, + {True: "key"}, + {1, 2}, + object(), + datetime(2026, 1, 1), + float("nan"), + float("inf"), + -float("inf"), + ], +) +def test_non_json_values_are_rejected(value): + with pytest.raises(TypeError, match="json-serializable"): + validate_json({"nested": [value]}) + + +def test_cycles_are_rejected_but_shared_values_are_valid(): + value = [] + validate_json([value, value]) + value.append(value) + with pytest.raises(TypeError, match="circular reference"): + validate_json(value) + + +def test_nested_error_identifies_path_without_printing_value(): + with pytest.raises(TypeError, match=r"arguments\['items'\]\[0\]") as error: + validate_json({"items": [b"secret-content"]}, "arguments") + assert "secret-content" not in str(error.value) + + +def test_function_request_defaults_to_json(): + request = FunctionRequest(function_name="identity", args=[None], kwargs={}) + payload = request.to_input() + assert payload["serialization_format"] == "json" + assert FunctionRequest.from_input(payload) == request + + +@pytest.mark.parametrize("fmt", ["cloudpickle", "typo", "", None, False]) +def test_request_rejects_unsupported_formats(fmt): + with pytest.raises(ValueError, match="serialization_format"): + FunctionRequest.from_input( + {"function_name": "identity", "serialization_format": fmt} + ) + + +@pytest.mark.parametrize("args,kwargs", [(None, {}), ({}, {}), ([], []), ([], None)]) +def test_request_rejects_invalid_argument_containers(args, kwargs): + with pytest.raises(TypeError, match="args must be a list"): + FunctionRequest.from_input( + {"function_name": "identity", "args": args, "kwargs": kwargs} + ) + + +def test_request_revalidates_mutations_before_encoding(): + request = FunctionRequest(function_name="identity") + request.args.append(b"bytes") + with pytest.raises(TypeError, match="json-serializable"): + request.to_input() + + +@pytest.mark.parametrize("kind", ["live", "deployed", "task"]) +@pytest.mark.parametrize("value", [b"bytes", (1, 2), {1: "key"}, float("nan")]) +def test_all_targets_reject_non_json_arguments(kind, value, monkeypatch): + spec = ResourceSpec( + kind=ResourceKind.TASK if kind == "task" else ResourceKind.QUEUE, + name="identity", + ) + target = { + "live": LiveTarget("ep"), + "deployed": SentinelTarget("a", "e", "q"), + "task": PodTarget(spec, identity), + }[kind] + monkeypatch.setattr( + "runpod.apps.serialization.get_function_source", + lambda fn: "def identity(value): return value", + ) + with pytest.raises(TypeError, match="identity.*json-serializable"): + target.build_payload(identity, spec, (value,), {}) + + +def test_task_and_live_requests_share_json_payloads(monkeypatch): + spec = ResourceSpec(kind=ResourceKind.TASK, name="identity") + monkeypatch.setattr( + "runpod.apps.serialization.get_function_source", + lambda fn: "def identity(value): return value", + ) + args = ({"items": [1, None]},) + task = PodTarget(spec, identity).build_payload(identity, spec, args, {}) + live = LiveTarget("ep").build_payload(identity, spec, args, {})["input"] + assert task == live + assert task["args"] == list(args) + assert task["serialization_format"] == "json" + + +@pytest.mark.parametrize("value", [None, False, 0, "", [], {}, {"x": [1, 2]}]) +def test_task_and_live_decode_identically(value): + response = {"success": True, "json_result": value} + assert unwrap_task_response(response) == value + assert LiveTarget("ep")._unwrap_response(response) == value + + +@pytest.mark.parametrize( + "response", + [ + {"success": True}, + {"success": True, "result": "encoded-pickle"}, + {"success": True, "result": "encoded-pickle", "json_result": 42}, + {"success": True, "serialization_format": "cloudpickle", "json_result": 42}, + {"success": True, "json_result": float("nan")}, + {"success": "true", "json_result": 42}, + ], +) +def test_task_and_live_reject_invalid_envelopes(response): + for decode in (unwrap_task_response, LiveTarget("ep")._unwrap_response): + with pytest.raises(RemoteExecutionError, match="invalid remote response"): + decode(response) + + +def test_response_requires_boolean_success(): + with pytest.raises(TypeError, match="boolean success"): + FunctionResponse.from_output({"json_result": 1}) + + +@pytest.fixture +def app(monkeypatch): + _clear_registry() + for variable in ("RUNPOD_ENDPOINT_ID", "RUNPOD_POD_ID", "RUNPOD_DEV_SESSION"): + monkeypatch.delenv(variable, raising=False) + yield App("json-contract") + _clear_registry() + + +@pytest.mark.parametrize("kind", ["queue", "task"]) +async def test_invalid_arguments_fail_before_resource_resolution(app, kind): + handle = getattr(app, kind)(name="identity")(identity) + resolve = AsyncMock() + app._resolve = resolve + for call in (handle.remote.aio, handle.spawn.aio): + with pytest.raises(TypeError, match="json-serializable"): + await call(b"bytes") + resolve.assert_not_called() + + +@pytest.mark.parametrize("kind", ["queue", "task"]) +def test_same_worker_remote_obeys_json_but_local_is_python(app, kind, monkeypatch): + handle = getattr(app, kind)(name="identity")(identity) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setattr(handle, "_is_current_worker", lambda: True) + with pytest.raises(TypeError, match="json-serializable"): + handle.remote(b"bytes") + assert handle.local(b"bytes") == b"bytes" + assert handle.remote({"x": [1, None]}) == {"x": [1, None]} + + +@pytest.mark.parametrize("asynchronous", [True, False]) +async def test_same_worker_validates_results(app, monkeypatch, asynchronous): + async def async_value(): + return b"bytes" + + def sync_value(): + return b"bytes" + + handle = app.queue(name="value")(async_value if asynchronous else sync_value) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setattr(handle, "_is_current_worker", lambda: True) + with pytest.raises(TypeError, match="return value.*json-serializable"): + await handle.remote.aio() + + +@pytest.mark.parametrize("asynchronous", [True, False]) +async def test_same_worker_validates_stream_chunks(app, monkeypatch, asynchronous): + async def async_stream(): + yield 1 + yield b"bytes" + + def sync_stream(): + yield 1 + yield b"bytes" + + handle = app.queue(name="stream")(async_stream if asynchronous else sync_stream) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setattr(handle, "_is_current_worker", lambda: True) + stream = handle.stream.aio() + assert await anext(stream) == 1 + with pytest.raises(TypeError, match="stream chunk.*json-serializable"): + await anext(stream) + + +@pytest.mark.parametrize("asynchronous", [True, False]) +async def test_same_worker_aggregate_stops_at_invalid_chunk( + app, monkeypatch, asynchronous +): + advanced = [] + + async def async_stream(): + yield (1, 2) + advanced.append(True) + + def sync_stream(): + yield (1, 2) + advanced.append(True) + + handle = app.queue(name="stream")(async_stream if asynchronous else sync_stream) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep") + monkeypatch.setattr(handle, "_is_current_worker", lambda: True) + with pytest.raises(TypeError, match="return value.*json-serializable"): + await handle.remote.aio() + assert advanced == [] + + +@pytest.mark.parametrize( + "value", [b"bytes", datetime(2026, 1, 1), float("nan"), (1, 2)] +) +async def test_serverless_reports_non_json_results(value): + result = await rp_job.run_job(lambda job: value, {"id": "test", "input": {}}) + error = json.loads(result["error"]) + assert "json-serializable" in error["error_message"] + assert "output" not in result + + +async def test_serverless_rejects_input_before_handler(): + called = [] + result = await rp_job.run_job( + lambda job: called.append(job), {"id": "test", "input": {"x": float("nan")}} + ) + assert not called + assert "job input" in json.loads(result["error"])["error_message"] + + +@pytest.mark.parametrize("asynchronous", [True, False]) +async def test_serverless_stream_reports_invalid_chunk(asynchronous): + async def async_stream(job): + yield 1 + yield b"bytes" + + def sync_stream(job): + yield 1 + yield b"bytes" + + handler = async_stream if asynchronous else sync_stream + chunks = [item async for item in rp_job.run_job_generator(handler, {"id": "test"})] + assert chunks[0] == {"output": 1} + assert "json-serializable" in chunks[1]["error"] + assert len(chunks) == 2 + + +@pytest.mark.parametrize("stream", [True, False]) +async def test_http_serialization_failure_sends_terminal_error(monkeypatch, stream): + transmit = AsyncMock() + monkeypatch.setattr(rp_http, "_transmit", transmit) + session = SimpleNamespace(headers={}) + send = rp_http.stream_result if stream else rp_http.send_result + await send(session, {"output": b"bytes"}, {"id": "test"}) + transmit.assert_awaited_once() + _, url, body = transmit.await_args.args + assert url == rp_http.JOB_DONE_URL.replace("$ID", "test") + "&isStream=false" + error = json.loads(json.loads(body)["error"]) + assert error["error_type"] == "SerializationError" + assert "json-serializable" in error["error_message"] + + +@pytest.mark.parametrize("route", ["runsync", "status", "stream"]) +async def test_local_api_reports_invalid_stream_chunk(monkeypatch, route): + def handler(job): + yield "first" + yield datetime(2026, 1, 1) + + removed = [] + monkeypatch.setattr( + rp_fastapi, + "job_list", + SimpleNamespace( + get=lambda job_id: SimpleNamespace(id=job_id, input={}, webhook=None), + remove=removed.append, + ), + ) + worker = object.__new__(rp_fastapi.WorkerAPI) + worker.config = {"handler": handler} + argument = rp_fastapi.DefaultRequest(input={}) if route == "runsync" else "test" + result = await getattr(worker, f"_sim_{route}")(argument) + assert result["status"] == "FAILED" + assert "json-serializable" in result["error"] + assert "output" not in result + if route != "runsync": + assert removed == ["test"] + + +@pytest.mark.parametrize("value", [None, {}, {"items": []}]) +async def test_local_api_preserves_json_results(value): + worker = object.__new__(rp_fastapi.WorkerAPI) + worker.config = {"handler": lambda job: value} + result = await worker._sim_runsync(rp_fastapi.DefaultRequest(input={})) + assert result["status"] == "COMPLETED" + assert result["output"] == value + + +@pytest.mark.parametrize("value", [None, {}, [], False, 0, ""]) +async def test_serverless_preserves_empty_json_results(value): + result = await rp_job.run_job(lambda job: value, {"id": "test", "input": {}}) + assert result == {"output": value} + + +@pytest.mark.parametrize("route", ["runsync", "status/test"]) +@pytest.mark.parametrize("value", [None, {}]) +def test_local_http_preserves_empty_json_results(monkeypatch, route, value): + from fastapi.testclient import TestClient + + monkeypatch.setattr(rp_fastapi.Heartbeat, "start_ping", lambda self, mirror: None) + monkeypatch.setattr(rp_fastapi.rp_capture, "install", lambda **kwargs: None) + monkeypatch.setattr( + rp_fastapi, + "job_list", + SimpleNamespace( + set_mirror=lambda mirror: None, + get=lambda job_id: SimpleNamespace(id=job_id, input={}, webhook=None), + remove=lambda job_id: None, + ), + ) + worker = rp_fastapi.WorkerAPI({"handler": lambda job: value}) + client = TestClient(worker.rp_app) + try: + response = client.post(f"/{route}", json={"input": {}}) + finally: + client.close() + assert response.status_code == 200 + assert response.json()["status"] == "COMPLETED" + assert "output" in response.json() + assert response.json()["output"] == value + + +@pytest.mark.parametrize("stream", [True, False]) +async def test_deployed_job_sends_serialization_failure(monkeypatch, stream): + def generator(job): + yield "first" + yield b"bytes" + + send = AsyncMock() + send_chunk = AsyncMock() + monkeypatch.setattr(rp_job, "send_result", send) + monkeypatch.setattr(rp_job, "stream_result", send_chunk) + handler = generator if stream else lambda job: datetime(2026, 1, 1) + await rp_job.handle_job( + SimpleNamespace(headers={}), + {"handler": handler, "return_aggregate_stream": True}, + {"id": "test", "input": {}}, + ) + send.assert_awaited_once() + result = send.await_args.args[1] + assert "json-serializable" in result["error"] + assert "output" not in result + assert send_chunk.await_count == (1 if stream else 0)