diff --git a/CHANGELOG.md b/CHANGELOG.md index dd2e5e695..dca322ac6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,13 @@ ### Features * Apps expose `NetworkVolume` and `GlobalVolume` resources with explicit path-keyed mounts, lazy create-if-missing provisioning, and execution-bound filesystem access. Global volumes use GraphQL for lookup and creation. +* Deployed Flash calls use endpoint or pod runtime credentials. Release requires deployed support from [RunPod#6686](https://github.com/runpod/RunPod/pull/6686), [ai-api#1176](https://github.com/runpod/ai-api/pull/1176), and [sdk-runtime#4](https://github.com/runpod-workers/sdk-runtime/pull/4). + +### Breaking Changes + +* Production endpoint environment variables reserve `RUNPOD_API_KEY` and runtime identity fields. Workers authenticate Flash calls with platform-issued runtime credentials; application credentials must use non-reserved environment names. +* Deployed workers can invoke Flash resources only within their own app and environment, including through `Queue`, `Api`, and `Task` stubs. Local clients use account credentials subject to their permissions. +* Remote queues, tasks, and JSON API calls require finite JSON-compatible arguments and results in local, dev, and deployed execution. Dev workers receive source separately from the JSON argument contract. Explicit `.local()` calls run Python directly. ### Bug Fixes diff --git a/README.md b/README.md index 6f5a3fc4d..b3fa882c0 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,18 @@ make setup make test ``` +## Remote values + +Queue and task calls use finite JSON-compatible arguments, return values, +and stream chunks. JSON conversion turns tuples into lists and supported +dictionary keys into strings, including in same-worker remote calls. +Bytes, custom objects, circular references, and non-finite numbers produce +explicit errors. `.local()` remains a direct Python call. + +Successful function responses contain `json_result`, including for a null +return value. SDK and worker runtimes share the JSON protocol validators; +use compatible package releases 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/__init__.py b/runpod/__init__.py index 427023537..6add070c1 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -54,6 +54,7 @@ "Model": "runpod.apps", "NetworkVolume": "runpod.apps", "Queue": "runpod.apps", + "Task": "runpod.apps", "Secret": "runpod.apps", "Volume": "runpod.apps", "is_local": "runpod.apps", diff --git a/runpod/_serialization.py b/runpod/_serialization.py new file mode 100644 index 000000000..7d4beb7e9 --- /dev/null +++ b/runpod/_serialization.py @@ -0,0 +1,21 @@ +"""finite json validation and wire-value conversion for callers and workers.""" + +import json +from typing import Any + + +def _json_text(value: Any, label: str) -> str: + try: + return json.dumps(value, allow_nan=False) + except (TypeError, ValueError, RecursionError) as exc: + raise TypeError(f"{label} must be json-serializable: {exc}") from exc + + +def validate_json(value: Any, label: str = "value") -> None: + """require values encodable as finite json.""" + _json_text(value, label) + + +def normalize_json(value: Any, label: str = "value") -> Any: + """apply json's container and key conversions without a network roundtrip.""" + return json.loads(_json_text(value, label)) diff --git a/runpod/api/graphql.py b/runpod/api/graphql.py index 9eeddaf33..48bd264db 100644 --- a/runpod/api/graphql.py +++ b/runpod/api/graphql.py @@ -87,6 +87,7 @@ async def run_graphql_query_async( variables: Optional[Dict[str, Any]] = None, timeout: float = 60.0, anonymous: bool = False, + headers: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """ Async variant of run_graphql_query, sharing the same url, headers, @@ -109,7 +110,7 @@ async def run_graphql_query_async( async with aiohttp.ClientSession(timeout=client_timeout) as session: async with session.post( _graphql_url(), - headers=_build_headers(effective_api_key), + headers={**_build_headers(effective_api_key), **(headers or {})}, json=_build_payload(query, variables), ) as response: if response.status == HTTP_STATUS_UNAUTHORIZED: diff --git a/runpod/api/mutations/apps.py b/runpod/api/mutations/apps.py index 2e3c73a44..1b549580e 100644 --- a/runpod/api/mutations/apps.py +++ b/runpod/api/mutations/apps.py @@ -92,6 +92,12 @@ } """ +MUTATION_FINALIZE_TASK_ARTIFACT_UPLOAD = """ +mutation FinalizeArtifactUpload($input: FinalizeFlashArtifactUploadInput!) { + finalizeFlashArtifactUpload(input: $input) { id manifest resourceSpec } +} +""" + MUTATION_DEPLOY_BUILD = """ mutation deployBuildToEnvironment($input: DeployBuildToEnvironmentInput!) { deployBuildToEnvironment(input: $input) { id name } diff --git a/runpod/apps/__init__.py b/runpod/apps/__init__.py index 6a62fef93..ceb449a49 100644 --- a/runpod/apps/__init__.py +++ b/runpod/apps/__init__.py @@ -15,7 +15,7 @@ from .markers import delete, get, init, patch, post, put from .schedule import schedule from .spec import ResourceKind, ResourceSpec, RouteSpec -from .stubs import Api, Queue +from .stubs import Api, Queue, Task from .model import Model from .secret import Secret from .volume import GlobalVolume, NetworkVolume, Volume @@ -34,6 +34,7 @@ "Model", "NetworkVolume", "Queue", + "Task", "Secret", "Volume", "RemoteExecutionError", diff --git a/runpod/apps/api.py b/runpod/apps/api.py index daca51308..c6195c398 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -20,6 +20,46 @@ _TRANSPORT_RETRIES = 4 +class FlashTaskApiClient: + """named task lifecycle operations with runtime-scoped authentication.""" + + async def _execute(self, query: str, task_input: Dict[str, Any]) -> Dict[str, Any]: + from .utils.network import api_key, runtime_headers + + headers = runtime_headers() + response = await run_graphql_query_async( + query, + variables={"input": task_input}, + api_key=(headers["Authorization"][7:] if headers else api_key()), + headers=headers, + ) + return response["data"] + + async def create(self, task_input: Dict[str, Any]) -> Dict[str, Any]: + from .utils.network import validate_runtime_target + + validate_runtime_target( + task_input.get("flashAppName"), task_input.get("flashEnvironmentName") + ) + data = await self._execute( + """mutation FlashTaskCreate($input: FlashTaskCreateInput!) { + flashTaskCreate(input: $input) { podId flashBuildId resourceName } + }""", + task_input, + ) + return data["flashTaskCreate"] + + async def terminate(self, pod_id: str, token: str) -> None: + data = await self._execute( + """mutation FlashTaskTerminate($input: FlashTaskTerminateInput!) { + flashTaskTerminate(input: $input) + }""", + {"podId": pod_id, "executionToken": token}, + ) + if data.get("flashTaskTerminate") is not True: + raise RuntimeError("task termination was not acknowledged") + + def _rest_endpoint_input(endpoint: Dict[str, Any]) -> Optional[Dict[str, Any]]: """translate only saves whose capabilities rest can preserve atomically.""" fields = { @@ -583,6 +623,8 @@ async def finalize_artifact_upload( self, app_id: str, object_key: str, manifest: Dict[str, Any] ) -> Dict[str, Any]: mutation = app_mutations.MUTATION_FINALIZE_ARTIFACT_UPLOAD + if any(r.get("kind") == "task" for r in manifest.get("resources", [])): + mutation = app_mutations.MUTATION_FINALIZE_TASK_ARTIFACT_UPLOAD data = await self._execute( mutation, { diff --git a/runpod/apps/app.py b/runpod/apps/app.py index c2a6a0e5d..b64855559 100644 --- a/runpod/apps/app.py +++ b/runpod/apps/app.py @@ -151,7 +151,11 @@ def task( accelerate_downloads: bool = True, container_disk_gb: Optional[int] = None, ) -> Callable[[Callable], FunctionHandle]: - """declare ephemeral pod compute from a function.""" + """declare ephemeral pod compute from a function. + + remote task arguments and results use json in every execution + context. use .local() for direct python execution without transport. + """ def decorator(fn: Callable) -> FunctionHandle: spec = ResourceSpec( @@ -259,6 +263,12 @@ async def _resolve(self, spec: ResourceSpec) -> InvocationTarget: """resolve a resource spec to an invocation target.""" spec.validate() if spec.kind is ResourceKind.TASK: + from .context import deployed_identity + from .targets import DeployedTaskTarget + + identity = deployed_identity() + if identity is not None: + return DeployedTaskTarget(self.name, identity[1], spec.name) handle = self._resources.get(spec.name) fn = getattr(handle, "_fn", None) return PodTarget( 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/context.py b/runpod/apps/context.py index b21d38760..46283fbce 100644 --- a/runpod/apps/context.py +++ b/runpod/apps/context.py @@ -7,7 +7,7 @@ import threading import time from enum import Enum -from typing import Any, Coroutine +from typing import Any, Coroutine, Optional, Tuple log = logging.getLogger(__name__) CLEANUP_TIMEOUT = 30.0 @@ -36,6 +36,27 @@ def current_context() -> Context: return Context.LOCAL +def deployed_identity() -> Optional[Tuple[str, str, str]]: + """app, environment and resource identity for a deployed worker.""" + if not ( + os.getenv("RUNPOD_RUNTIME_KIND") == "task" + or os.getenv("RUNPOD_BUILD_ID") + or os.getenv("FLASH_APP_NAME") + or os.getenv("FLASH_RESOURCE_NAME") + ): + return None + fields = tuple( + os.getenv(key) + for key in ("FLASH_APP_NAME", "FLASH_ENVIRONMENT", "FLASH_RESOURCE_NAME") + ) + build = os.getenv("RUNPOD_BUILD_ID") + if any(not value or value != value.strip() for value in (*fields, build)): + raise RuntimeError( + "deployed flash requires app, environment, resource and build identity" + ) + return fields + + def is_local() -> bool: """true when not running inside a runpod container. diff --git a/runpod/apps/deploy.py b/runpod/apps/deploy.py index bddd63615..bd57ba639 100644 --- a/runpod/apps/deploy.py +++ b/runpod/apps/deploy.py @@ -6,13 +6,11 @@ import tarfile import tempfile from dataclasses import dataclass, field -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Any, Dict, List, Optional, Set, Tuple from pathspec import GitIgnoreSpec -import runpod - from .api import AppsApiClient from .app import App from .build import ( @@ -115,12 +113,36 @@ class DeployResult: endpoints: Dict[str, str] = field(default_factory=dict) +def _validate_endpoint_env(env: Dict[str, str]) -> None: + reserved = { + "RUNPOD_API_KEY", + "RUNPOD_AI_API_KEY", + "RUNPOD_ENDPOINT_ID", + "RUNPOD_POD_ID", + "RUNPOD_TASK_TOKEN", + "RUNPOD_TASK_PORT", + "FLASH_BUILD_ARTIFACT_PATH", + "RUNPOD_RUNTIME_KIND", + "RUNPOD_DEV_APP", + "RUNPOD_DEV_RESOURCE", + "RUNPOD_DEV_SESSION", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_BUILD_ID", + } + if conflicts := reserved.intersection(env): + raise InvalidResourceError( + "production endpoint env contains reserved runtime fields: " + + ", ".join(sorted(conflicts)) + ) + + def build_manifest( app: App, project_root: Path, *, python_version: str = DEFAULT_PYTHON_VERSION, - excluded_packages: Optional[List[str]] = None, ) -> Dict[str, Any]: """serialize an app's resources for the coordinator and workers.""" resources = [] @@ -137,6 +159,48 @@ def build_manifest( ) entry = spec.to_manifest() + if spec.kind in (ResourceKind.QUEUE, ResourceKind.API): + _validate_endpoint_env(entry.get("env", {})) + if spec.kind is ResourceKind.TASK: + from .images import image_for_spec + + for key in entry.get("env", {}): + normalized = key.strip().upper() + if normalized.startswith(("RUNPOD_", "FLASH_")) or normalized in { + "PYTHONPATH", + "PYTHONHOME", + "LD_PRELOAD", + }: + raise InvalidResourceError( + f"task '{spec.name}' env contains reserved runtime field {key!r}" + ) + for path in spec.mounts: + mount_path = PurePosixPath(path) + for reserved in ( + "/root/.runpod", + "/runpod-volume", + "/proc", + "/sys", + "/dev", + ): + reserved_path = PurePosixPath(reserved) + if ( + mount_path == reserved_path + or mount_path in reserved_path.parents + or reserved_path in mount_path.parents + ): + raise InvalidResourceError( + f"task '{spec.name}' mount {path!r} overlaps " + f"reserved runtime path {reserved!r}" + ) + entry["imageName"] = image_for_spec(spec, python_version=python_version) + entry["customImage"] = bool(spec.image) + for mount, volume in zip(entry.get("mounts", []), spec.mounts.values()): + mount["create"] = volume.create + if volume.kind == "network": + mount["size"] = volume.size + if volume.datacenter is not None: + mount["datacenter"] = volume.datacenter fn = ( getattr(handle, "_fn", None) @@ -156,12 +220,34 @@ def build_manifest( "resources": resources, } - if excluded_packages: - manifest["excludedPackages"] = excluded_packages - return manifest +async def _prepare_task_resources( + app: App, + manifest: Dict[str, Any], + client: AppsApiClient, + events=None, + *, + resolver=None, +) -> None: + from .registry import resolve_registry_auth + from .volume import VolumeResolver + + resolver = resolver or VolumeResolver(client, events=events) + specs = [handle.spec for handle in app.resources.values()] + for resource in manifest["resources"]: + if resource["kind"] != "task": + continue + spec = app.resources[resource["name"]].spec + bindings = await resolver.resolve_mounts(spec.mounts, specs) + for mount, binding in zip(resource.get("mounts", []), bindings): + mount["id"] = binding["id"] + registry_id = await resolve_registry_auth(spec.registry_auth, api=client) + if registry_id: + resource["registryAuth"] = registry_id + + def _load_ignores(path: Path) -> GitIgnoreSpec: if path.is_symlink() or not path.is_file(): return GitIgnoreSpec.from_lines([]) @@ -300,21 +386,16 @@ def _deployed_endpoint_input( from .secret import render_env spec.validate() + resource_env = render_env(spec.env) + _validate_endpoint_env(resource_env) template_env = { + **resource_env, + "FLASH_APP_NAME": app.name, + "FLASH_ENVIRONMENT": env_name or app.env, "FLASH_RESOURCE_NAME": spec.name, - # version-triggering: a new build recreates all workers "RUNPOD_BUILD_ID": build_id, - **render_env(spec.env), - "FLASH_ENVIRONMENT": env_name or app.env, } - # cross-resource calls from inside a worker go through the sentinel, - # which authenticates with the account api key; without it any - # .remote() from worker code would fail - api_key = runpod.api_key - if api_key and "RUNPOD_API_KEY" not in template_env: - template_env["RUNPOD_API_KEY"] = api_key - # nested tasks use the same runtime and sdk channels as the endpoint for key in ( "RUNPOD_RUNTIME_TAG", @@ -391,6 +472,8 @@ async def reconcile_endpoints( environment: Dict[str, Any], build_id: str, python_version: str, + *, + resolver=None, ) -> Dict[str, str]: """converge the environment's endpoints to the app's resources. @@ -411,7 +494,7 @@ async def reconcile_endpoints( from .registry import resolve_registry_auth from .volume import VolumeResolver, attach_endpoint_volumes - resolver = VolumeResolver(client) + resolver = resolver or VolumeResolver(client) endpoints: Dict[str, str] = {} for name, handle in sorted(provisionable.items()): members = None @@ -467,6 +550,7 @@ def build_artifact( output: Optional[Path] = None, ) -> Path: """vendor, build the manifest, and package one app's artifact.""" + manifest = build_manifest(app, project_root, python_version=python_version) _phase(events, "vendor", f"python {python_version}") build: BuildResult = build_environment( app, @@ -475,12 +559,8 @@ def build_artifact( exclude=exclude, events=events, ) - manifest = build_manifest( - app, - project_root, - python_version=python_version, - excluded_packages=build.excluded, - ) + if build.excluded: + manifest["excludedPackages"] = build.excluded _phase(events, "package") tar_path = package_project( project_root, manifest, output=output, env_dir=build.env_dir @@ -512,6 +592,8 @@ async def deploy_app( for handle in app.resources.values(): handle.spec.validate() + manifest = build_manifest(app, project_root, python_version=python_version) + # fail fast on unresolvable secret references (workers would boot # with the literal template string otherwise) from .secret import secret_names, validate_secrets @@ -531,12 +613,12 @@ async def deploy_app( exclude=exclude, events=events, ) - manifest = build_manifest( - app, - project_root, - python_version=python_version, - excluded_packages=build.excluded, - ) + if build.excluded: + manifest["excludedPackages"] = build.excluded + from .volume import VolumeResolver + + resolver = VolumeResolver(client, events=events) + await _prepare_task_resources(app, manifest, client, events, resolver=resolver) _phase(events, "package") tar_path = package_project(project_root, manifest, env_dir=build.env_dir) tar_size = tar_path.stat().st_size @@ -568,6 +650,20 @@ async def deploy_app( upload["uploadUrl"], str(tar_path), progress=upload_progress ) build = await client.finalize_artifact_upload(app_id, upload["objectKey"], manifest) + definitions = build.get("resourceSpec") + tasks = {r["name"] for r in manifest["resources"] if r["kind"] == "task"} + if tasks and ( + not isinstance(definitions, dict) + or not isinstance(definitions.get("tasks"), dict) + or set(definitions["tasks"]) != tasks + or any( + not isinstance(task, dict) or not task + for task in definitions["tasks"].values() + ) + ): + raise RuntimeError( + "backend did not acknowledge the registered task definitions" + ) log.info("uploaded build %s", build["id"]) await client.deploy_build(environment["id"], build["id"]) @@ -575,7 +671,7 @@ async def deploy_app( _phase(events, "endpoints") endpoints = await reconcile_endpoints( - client, app, environment, build["id"], python_version + client, app, environment, build["id"], python_version, resolver=resolver ) endpoint_ready = getattr(events, "endpoint_ready", None) if endpoint_ready is not None: diff --git a/runpod/apps/handles.py b/runpod/apps/handles.py index b66af1d94..e4e804f61 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 normalize_json from .context import Context, block, current_context from .discovery_state import DiscoveryInvocationError, in_discovery from .errors import InvalidResourceError @@ -122,6 +123,15 @@ def _call_local_sync(self, *args: Any, **kwargs: Any) -> Any: result = self._fn(*args, **kwargs) return list(result) if inspect.isgenerator(result) else result + @staticmethod + def _json_value(value: Any) -> Any: + try: + return normalize_json(value) + except TypeError: + raise ValueError( + "remote arguments and results must be JSON-serializable" + ) from None + @staticmethod async def _await_local_result(result: Any) -> Any: if inspect.isawaitable(result): @@ -139,10 +149,11 @@ def _remote_sync(self, *args: Any, **kwargs: Any) -> Any: ): self._guard_discovery() self.spec.validate() + args, kwargs = self._json_value([args, kwargs]) result = self._call_local_sync(*args, **kwargs) if inspect.isawaitable(result) or inspect.isasyncgen(result): - return block(self._await_local_result(result)) - return result + result = block(self._await_local_result(result)) + return self._json_value(result) return block(self._remote_async(*args, **kwargs)) async def _remote_async(self, *args: Any, **kwargs: Any) -> Any: @@ -151,13 +162,14 @@ async def _remote_async(self, *args: Any, **kwargs: Any) -> Any: ctx = current_context() if ctx is Context.WORKER and self._is_current_worker(): + args, kwargs = self._json_value([args, kwargs]) if inspect.iscoroutinefunction(self._fn) or inspect.isasyncgenfunction( self._fn ): result = self._fn(*args, **kwargs) else: result = await asyncio.to_thread(self._call_local_sync, *args, **kwargs) - return await self._await_local_result(result) + return self._json_value(await self._await_local_result(result)) target = await self._app._resolve(self.spec) payload = target.build_payload(self._fn, self.spec, args, kwargs) @@ -172,13 +184,14 @@ async def _stream_async(self, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: ctx = current_context() if ctx is Context.WORKER and self._is_current_worker(): + args, kwargs = self._json_value([args, kwargs]) gen = self._fn(*args, **kwargs) if inspect.isasyncgen(gen): async for chunk in gen: - yield chunk + yield self._json_value(chunk) else: for chunk in gen: - yield chunk + yield self._json_value(chunk) return target = await self._app._resolve(self.spec) @@ -233,6 +246,15 @@ def _guard_discovery(self) -> None: def _is_current_worker(self) -> bool: import os + from .context import deployed_identity + + identity = deployed_identity() + if identity is not None: + return ( + self.spec.kind is not ResourceKind.TASK + and identity[0] == self._app.name + and identity[2] == self.spec.name + ) if dev_app := os.getenv("RUNPOD_DEV_APP"): return ( dev_app == self._app.name 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..dea7d5cfa 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 normalize_json, validate_json + EXECUTION_FUNCTION = "function" -FORMAT_CLOUDPICKLE = "cloudpickle" FORMAT_JSON = "json" @@ -31,10 +30,30 @@ 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() + request.args = normalize_json(request.args, "arguments") + request.kwargs = normalize_json(request.kwargs, "keyword arguments") + 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,17 +63,24 @@ 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"), - json_result=output.get("json_result"), + success=output["success"], + json_result=normalize_json(output.get("json_result"), "return value"), error=output.get("error"), stdout=output.get("stdout"), ) diff --git a/runpod/apps/serialization.py b/runpod/apps/serialization.py index 425189e1a..0972e2859 100644 --- a/runpod/apps/serialization.py +++ b/runpod/apps/serialization.py @@ -1,44 +1,38 @@ """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. +remote calls use finite json-compatible arguments and results. this +module validates that contract, decodes json response envelopes, and extracts +function source for live workers. """ import ast -import base64 import inspect import os import textwrap -from typing import Any, Callable, Dict, List +from typing import Any, Callable -import cloudpickle +from .._serialization import validate_json +from .errors import RemoteExecutionError +from .protocol import FunctionResponse -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] - - -def serialize_kwargs(kwargs: dict) -> Dict[str, str]: - return {k: serialize_arg(v) for k, v in kwargs.items()} +def unwrap_function_response(output: dict) -> Any: + try: + 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"remote execution failed: {response.error or 'unknown'}" + ) + return response.json_result -def deserialize_result(result_b64: str) -> Any: +def require_json_value(value: Any, message: str) -> None: try: - return cloudpickle.loads(base64.b64decode(result_b64)) - except ModuleNotFoundError as exc: - from .errors import RemoteExecutionError - - 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 + validate_json(value) + except TypeError: + raise ValueError(message) from None def require_json_args(fn: Callable, args: tuple, kwargs: dict) -> None: @@ -48,16 +42,15 @@ def require_json_args(fn: Callable, args: tuple, kwargs: dict) -> None: 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: + require_json_value( + {"args": list(args), "kwargs": kwargs}, + "remote arguments must be json-serializable", + ) + except ValueError: raise TypeError( - f"{fn.__name__}() arguments must be json-serializable " - f"(queue jobs are plain json): {exc}" - ) from exc + f"{fn.__name__}() arguments must be json-serializable" + ) from None def get_function_source(fn: Callable) -> str: diff --git a/runpod/apps/stubs.py b/runpod/apps/stubs.py index f0e5aeef4..c091c2307 100644 --- a/runpod/apps/stubs.py +++ b/runpod/apps/stubs.py @@ -9,8 +9,10 @@ external_api = runpod.Api(app="other-app", name="inference") external_api.post("/generate", {"prompt": "hi"}) -resolution goes through the same server-side sentinel path as deployed -handles, so a stub is just a handle without a function body. +queue and api stubs resolve through the sentinel; task stubs use the +named task api. local clients use account credentials, subject to their +permissions. deployed workers use runtime credentials and can target +only their own app and environment, including when using stubs. """ import os @@ -123,6 +125,33 @@ def __repr__(self) -> str: return f"" +class Task(_StubBase): + """client for an approved task deployed in a flash environment.""" + + def __init__(self, *, app: str, name: str, env: Optional[str] = None): + super().__init__(app=app, name=name, env=env) + if not app or not name: + raise ValueError("task app and name must be non-empty") + self.remote = Invoker(self._remote_async) + self.spawn = Invoker(self._spawn_async) + + def _target(self): + from .targets import DeployedTaskTarget + + return DeployedTaskTarget(self.app_name, self.env, self.resource_name) + + async def _remote_async(self, *args: Any, **kwargs: Any) -> Any: + target = self._target() + return await target.invoke(target.build_payload(None, None, args, kwargs)) + + async def _spawn_async(self, *args: Any, **kwargs: Any) -> Any: + target = self._target() + return await target.submit(target.build_payload(None, None, args, kwargs)) + + def __repr__(self): + return f"" + + class _StubRouteCaller: __slots__ = ("_stub", "_method") diff --git a/runpod/apps/targets.py b/runpod/apps/targets.py index 28e720513..011641d53 100644 --- a/runpod/apps/targets.py +++ b/runpod/apps/targets.py @@ -15,7 +15,8 @@ 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 .serialization import require_json_value from .spec import ResourceSpec from .utils.events import emit from .utils.network import api_key as _api_key @@ -39,11 +40,16 @@ def _lb_domain() -> str: def _headers(extra: Optional[Dict[str, str]] = None) -> Dict[str, str]: + from .utils.network import runtime_headers + headers = { - "Authorization": f"Bearer {_api_key()}", "Content-Type": "application/json", "User-Agent": USER_AGENT, } + runtime = runtime_headers() + headers.update( + runtime if runtime is not None else {"Authorization": f"Bearer {_api_key()}"} + ) if extra: headers.update(extra) return headers @@ -117,6 +123,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.""" output = data.get("output", data) + require_json_value(output, "queue results must be JSON-serializable") if data.get("status") == "FAILED" or data.get("error"): err = data.get("error") if not err: @@ -148,6 +155,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 +228,7 @@ async def _request_json( retry_not_found: bool = False, ) -> Dict[str, Any]: """http json call with exponential backoff on transient failures.""" + require_json_value(payload, "remote request body must be JSON-serializable") retry_transient = method.upper() in {"GET", "HEAD", "OPTIONS"} client_timeout = aiohttp.ClientTimeout(total=timeout) @@ -452,6 +463,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 []: + require_json_value( + chunk.get("output"), "queue results must be JSON-serializable" + ) yield chunk.get("output") if data.get("status") in FINAL_STATUSES: return @@ -482,6 +496,7 @@ async def http( timeout: float = DEFAULT_TIMEOUT_SECONDS, ) -> Any: """plain http call through the endpoint's lb subdomain.""" + require_json_value(body, "api request body must be JSON-serializable") url = f"https://{self.endpoint_id}.{_lb_domain()}{path}" client_timeout = aiohttp.ClientTimeout(total=timeout) async with aiohttp.ClientSession(timeout=client_timeout) as session: @@ -493,9 +508,11 @@ async def http( resp.raise_for_status() text = await resp.text() try: - return json.loads(text) + result = json.loads(text) except json.JSONDecodeError: return text + require_json_value(result, "api results must be JSON-serializable") + return result class SentinelTarget(InvocationTarget): @@ -518,6 +535,9 @@ def __init__(self, app_name: str, env_name: str, resource_name: str): ) def _sentinel_headers(self) -> Dict[str, str]: + from .utils.network import validate_runtime_target + + validate_runtime_target(self.app_name, self.env_name) return _headers( { "X-Flash-App": self.app_name, @@ -674,17 +694,11 @@ def unwrap(self, data: Dict[str, Any]) -> Any: return self._unwrap_response(output) def _unwrap_response(self, output: Any) -> Any: + require_json_value(output, "queue results must be JSON-serializable") 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) return output def _monitor(self): @@ -778,6 +792,7 @@ async def request( emit(self.events, "dispatch", self.resource_name, f"{method} {path}") start = time.monotonic() try: + require_json_value(body, "api request body must be JSON-serializable") await self._sync_source(timeout) result = await self._client.http(method, path, body, timeout=timeout) except Exception: @@ -797,63 +812,41 @@ async def request( return result -class PodTarget(InvocationTarget): - """ephemeral pod execution for @app.task: provision a pod, run the - function body, collect the result, terminate. - - tasks run to completion far beyond the queue data plane's sync - window, so the transport is a direct http channel to a single-shot - runner on the pod (see runpod.apps.tasks). - """ - - def __init__( - self, - spec: ResourceSpec, - fn: Callable, - events: Optional[object] = None, - *, - specs: Optional[List[ResourceSpec]] = None, +def _validate_task_payload(payload: Dict[str, Any]) -> None: + message = "tasks accept only JSON args and kwargs" + if ( + not isinstance(payload, dict) + or set(payload) != {"args", "kwargs"} + or not isinstance(payload["args"], list) + or not isinstance(payload["kwargs"], dict) + or any(not isinstance(key, str) for key in payload["kwargs"]) ): - self.spec = spec - self.fn = fn - self.events = events - self.specs = specs + raise ValueError(message) + require_json_value(payload, message) - 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, - ) - request = FunctionRequest( - function_name=fn.__name__, - function_code=get_function_source(fn), - args=serialize_args(args), - kwargs=serialize_kwargs(kwargs), - dependencies=spec.dependencies, - system_dependencies=spec.system_dependencies, - accelerate_downloads=spec.accelerate_downloads, - ) - return request.to_input() +class _TaskTarget(InvocationTarget): + """shared execution, result collection, and cleanup for task targets.""" + + @abstractmethod + def _execution(self, payload, timeout=None): + raise NotImplementedError async def invoke( self, payload: Dict[str, Any], *, timeout: Optional[float] = None ) -> Any: import time - from .tasks import TaskExecution, finish_cleanup, unwrap_task_response + from .tasks import finish_cleanup, unwrap_task_response - name = self.spec.name - hardware = ",".join(self.spec.cpu or self.spec.gpu or ["any"]) + name = self.resource_name + hardware = self.hardware emit(self.events, "dispatch", name) emit(self.events, "task_status", name, f"provisioning pod on {hardware}") start = time.monotonic() stream = None - execution = TaskExecution(self.spec, specs=self.specs) + execution = self._execution(payload, timeout) try: await execution.start() if execution.pod_id: @@ -899,9 +892,9 @@ async def cleanup(): return result async def submit(self, payload: Dict[str, Any]) -> Any: - from .tasks import TaskExecution, TaskJob, finish_cleanup + from .tasks import TaskJob, finish_cleanup - execution = TaskExecution(self.spec, specs=self.specs) + execution = self._execution(payload) try: await execution.start() await execution.wait_ready() @@ -910,3 +903,73 @@ async def submit(self, payload: Dict[str, Any]) -> Any: await finish_cleanup(execution.terminate()) raise return TaskJob(execution) + + +class PodTarget(_TaskTarget): + """ephemeral pod execution for @app.task: provision a pod, run the + function body, collect the result, terminate. + + tasks run to completion far beyond the queue data plane's sync + window, so the transport is a direct http channel to a single-shot + runner on the pod (see runpod.apps.tasks). + """ + + def __init__( + self, + spec: ResourceSpec, + fn: Callable, + events: Optional[object] = None, + *, + specs: Optional[List[ResourceSpec]] = None, + ): + self.spec = spec + self.fn = fn + self.events = events + self.specs = specs + self.resource_name = spec.name + self.hardware = ",".join(spec.cpu or spec.gpu or ["any"]) + + def _execution(self, payload, timeout=None): + from .tasks import TaskExecution + + return TaskExecution(self.spec, specs=self.specs) + + def build_payload( + self, fn: Callable, spec: ResourceSpec, args: tuple, kwargs: dict + ) -> Dict[str, Any]: + from .serialization import get_function_source + + _validate_task_payload({"args": list(args), "kwargs": kwargs}) + request = FunctionRequest( + function_name=fn.__name__, + function_code=get_function_source(fn), + 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() + + +class DeployedTaskTarget(_TaskTarget): + """invoke an approved deployed task without transmitting source or specs.""" + + def __init__(self, app, env, name, events=None): + self.app_name = app + self.env_name = env + self.resource_name = name + self.hardware = "approved resources" + self.events = events + + def _execution(self, payload, timeout=None): + from .tasks import DeployedTaskExecution + + _validate_task_payload(payload) + return DeployedTaskExecution( + self.app_name, self.env_name, self.resource_name, timeout=timeout + ) + + def build_payload(self, fn, spec, args, kwargs): + return {"args": list(args), "kwargs": kwargs} diff --git a/runpod/apps/tasks.py b/runpod/apps/tasks.py index e6a6419b4..9663aee31 100644 --- a/runpod/apps/tasks.py +++ b/runpod/apps/tasks.py @@ -21,7 +21,6 @@ from .api import AppsApiClient, is_capacity_error from .context import CLEANUP_TIMEOUT -from .errors import RemoteExecutionError from .spec import ResourceSpec log = logging.getLogger(__name__) @@ -36,6 +35,16 @@ _pending_cleanups = set() +def _validate_timeout(timeout): + if timeout is not None and ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not math.isfinite(timeout) + or timeout < 0 + ): + raise ValueError("task timeout must be a finite non-negative number") + + async def finish_cleanup(coro) -> None: """finish cleanup despite repeated cancellation, without hiding its cause.""" original = sys.exc_info()[1] @@ -185,19 +194,11 @@ def _pod_input(spec: ResourceSpec, token: str, task_name: str) -> Dict[str, Any] return pod -class TaskExecution: - """one pod running one function.""" +class _TaskTransport: + """token-authenticated execution transport with cancellation-safe cleanup.""" - def __init__( - self, - spec: ResourceSpec, - api: Optional[AppsApiClient] = None, - *, - specs: Optional[List[ResourceSpec]] = None, - ): - self.spec = spec - self.specs = specs if specs is not None else [spec] - self.api = api or AppsApiClient() + def __init__(self, api): + self.api = api self.token = secrets.token_urlsafe(32) self.pod_id: Optional[str] = None self._deployment: Optional[asyncio.Task] = None @@ -206,55 +207,6 @@ def __init__( def _headers(self) -> Dict[str, str]: return {"Authorization": f"Bearer {self.token}"} - async def start(self) -> None: - pod = _pod_input(self.spec, self.token, self.spec.name) - if self.spec.registry_auth: - from .registry import resolve_registry_auth - - pod["containerRegistryAuthId"] = await resolve_registry_auth( - self.spec.registry_auth, api=self.api - ) - pod = await self._attach_mounts(pod) - self._deployment = asyncio.create_task(self._deploy_and_record(pod)) - await asyncio.shield(self._deployment) - - async def _deploy_and_record(self, pod: Dict[str, Any]) -> None: - result = await self._deploy_pod(pod) - self.pod_id = result["id"] - log.info("task pod %s deployed for %s", self.pod_id, self.spec.name) - - async def _deploy_pod(self, pod: Dict[str, Any]) -> Dict[str, Any]: - from ..error import QueryError - - try: - return await self.api.deploy_task_pod(pod, is_cpu=self.spec.is_cpu) - except QueryError as exc: - if not is_capacity_error(exc): - raise - capacity_error = exc - - datacenters = pod.get("dataCenterIds") or [] - if len(datacenters) < 2: - raise capacity_error - - for datacenter in datacenters: - candidate = dict(pod, dataCenterIds=[datacenter]) - try: - return await self.api.deploy_task_pod( - candidate, is_cpu=self.spec.is_cpu - ) - except QueryError as exc: - if not is_capacity_error(exc): - raise - raise capacity_error - - async def _attach_mounts(self, pod: Dict[str, Any]) -> Dict[str, Any]: - """resolve storage mounts and apply their placement constraints.""" - from .volume import VolumeResolver, attach_pod_mounts - - await attach_pod_mounts(pod, self.spec, VolumeResolver(self.api), self.specs) - return pod - async def wait_ready(self, timeout: float = READY_TIMEOUT) -> None: """poll the runner's /ping through the pod proxy until it answers.""" url = f"{_proxy_url(self.pod_id)}/ping" @@ -311,14 +263,11 @@ async def set_timeout(self, timeout: float) -> None: await self._post("/timeout", {"timeout": timeout}) async def _post(self, path: str, request: Dict[str, Any]) -> None: + from .serialization import require_json_value + timeout = request.get("timeout") - if timeout is not None and ( - isinstance(timeout, bool) - or not isinstance(timeout, (int, float)) - or not math.isfinite(timeout) - or timeout < 0 - ): - raise ValueError("task timeout must be a finite non-negative number") + _validate_timeout(timeout) + require_json_value(request, "task request body must be JSON-serializable") url = f"{_proxy_url(self.pod_id)}{path}" attempts = 6 async with aiohttp.ClientSession() as session: @@ -385,9 +334,7 @@ async def terminate(self) -> None: try: for attempt in range(DELETE_ATTEMPTS): try: - await asyncio.wait_for( - self.api.terminate_pod(self.pod_id), DELETE_TIMEOUT - ) + await asyncio.wait_for(self._terminate_pod(), DELETE_TIMEOUT) break except Exception as exc: status = getattr(exc, "status_code", None) or getattr( @@ -420,24 +367,124 @@ async def terminate(self) -> None: raise self.pod_id = None + async def _terminate_pod(self): + raise NotImplementedError + + +class TaskExecution(_TaskTransport): + """one pod running one function.""" + + def __init__( + self, + spec: ResourceSpec, + api: Optional[AppsApiClient] = None, + *, + specs: Optional[List[ResourceSpec]] = None, + ): + self.spec = spec + self.specs = specs if specs is not None else [spec] + super().__init__(api or AppsApiClient()) + + async def start(self) -> None: + pod = _pod_input(self.spec, self.token, self.spec.name) + if self.spec.registry_auth: + from .registry import resolve_registry_auth + + pod["containerRegistryAuthId"] = await resolve_registry_auth( + self.spec.registry_auth, api=self.api + ) + pod = await self._attach_mounts(pod) + self._deployment = asyncio.create_task(self._deploy_and_record(pod)) + await asyncio.shield(self._deployment) + + async def _deploy_and_record(self, pod: Dict[str, Any]) -> None: + result = await self._deploy_pod(pod) + self.pod_id = result["id"] + log.info("task pod %s deployed for %s", self.pod_id, self.spec.name) + + async def _deploy_pod(self, pod: Dict[str, Any]) -> Dict[str, Any]: + from ..error import QueryError + + try: + return await self.api.deploy_task_pod(pod, is_cpu=self.spec.is_cpu) + except QueryError as exc: + if not is_capacity_error(exc): + raise + capacity_error = exc + + datacenters = pod.get("dataCenterIds") or [] + if len(datacenters) < 2: + raise capacity_error + + for datacenter in datacenters: + candidate = dict(pod, dataCenterIds=[datacenter]) + try: + return await self.api.deploy_task_pod( + candidate, is_cpu=self.spec.is_cpu + ) + except QueryError as exc: + if not is_capacity_error(exc): + raise + raise capacity_error + + async def _attach_mounts(self, pod: Dict[str, Any]) -> Dict[str, Any]: + """resolve storage mounts and apply their placement constraints.""" + from .volume import VolumeResolver, attach_pod_mounts + + await attach_pod_mounts(pod, self.spec, VolumeResolver(self.api), self.specs) + return pod + + async def _terminate_pod(self): + await self.api.terminate_pod(self.pod_id) + + +class DeployedTaskExecution(_TaskTransport): + """one approved named task using the task transport and cleanup lifecycle.""" + + def __init__(self, app, env, name, api=None, *, timeout=None): + from .api import FlashTaskApiClient + + _validate_timeout(timeout) + if timeout == 0: + raise ValueError("deployed task timeout must be a finite positive number") + self.app_name = app + self.env_name = env + self.resource_name = name + self.timeout = timeout + super().__init__(api or FlashTaskApiClient()) + + async def start(self): + self._deployment = asyncio.create_task(self._create_and_record()) + await asyncio.shield(self._deployment) + + async def _create_and_record(self): + task_input = { + "flashAppName": self.app_name, + "flashEnvironmentName": self.env_name, + "resourceName": self.resource_name, + "executionToken": self.token, + } + if self.timeout is not None: + task_input["executionTimeoutSeconds"] = self.timeout + result = await self.api.create(task_input) + self.pod_id = result["podId"] + + async def _terminate_pod(self): + await self.api.terminate(self.pod_id, self.token) + 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 require_json_value, unwrap_function_response - return deserialize_result(response["result"]) - return response.get("json_result") + require_json_value(response, "task results must be JSON-serializable") + return unwrap_function_response(response) class TaskJob: """handle for a spawned task; owns the pod until the result is read.""" - def __init__(self, execution: TaskExecution): + def __init__(self, execution: _TaskTransport): self._execution = execution self._result: Any = None self._done = False diff --git a/runpod/apps/utils/network.py b/runpod/apps/utils/network.py index d681766ea..0b1818b59 100644 --- a/runpod/apps/utils/network.py +++ b/runpod/apps/utils/network.py @@ -1,11 +1,43 @@ """shared authentication and lightweight data-plane requests.""" import os -from typing import Any, Dict +from typing import Any, Dict, Optional import aiohttp +def runtime_headers() -> Optional[Dict[str, str]]: + from ..context import deployed_identity + + if deployed_identity() is None: + return None + endpoint = os.getenv("RUNPOD_ENDPOINT_ID") + pod = os.getenv("RUNPOD_POD_ID") + if os.getenv("RUNPOD_RUNTIME_KIND") == "task": + endpoint = None + if not endpoint and not pod: + raise RuntimeError("production flash runtime source identity is missing") + key = os.getenv("RUNPOD_AI_API_KEY" if endpoint else "RUNPOD_API_KEY") + if not key or key != key.strip(): + raise RuntimeError("production flash runtime credential is missing") + if (endpoint or pod) != (endpoint or pod).strip(): + raise RuntimeError("production flash runtime source identity is malformed") + return { + "Authorization": f"Bearer {key}", + "X-RunPod-Endpoint-Id" if endpoint else "X-RunPod-Pod-Id": endpoint or pod, + } + + +def validate_runtime_target(app_name: str, env_name: str) -> None: + from ..context import deployed_identity + + identity = deployed_identity() + if identity is not None and (app_name, env_name) != identity[:2]: + raise ValueError( + "deployed flash calls require the same app and environment as the worker" + ) + + def api_key() -> str: key = os.getenv("RUNPOD_API_KEY") if not key: diff --git a/runpod/rp_cli/console.py b/runpod/rp_cli/console.py index f165ff56e..fe0e48738 100644 --- a/runpod/rp_cli/console.py +++ b/runpod/rp_cli/console.py @@ -233,11 +233,12 @@ def resources_table(rows: Iterable[tuple]) -> None: w_name = max(len(r[0]) for r in rows) w_hw = max(len(r[2]) for r in rows) for name, kind, hardware, endpoint_id in rows: - tail = ( - "[dim]per-call[/dim]" - if not endpoint_id or endpoint_id == "per-call" - else endpoint_link(endpoint_id) - ) + if endpoint_id == "registered": + tail = "[dim]registered / per-call[/dim]" + elif not endpoint_id or endpoint_id == "per-call": + tail = "[dim]per-call[/dim]" + else: + tail = endpoint_link(endpoint_id) console.print( f" [white]{name:<{w_name}}[/white]" f" {kind_badge(kind)}" diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 8943d9960..e878f7bc1 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -287,14 +287,17 @@ def deploy(target, env, python_version, exclude, build_only): f"[bold white]{result.app_name}/{env or found.env}[/bold white] " f"is live [dim]{t.elapsed:.1f}s[/dim]" ) - if result.endpoints: + if result.resources: ui.console.print() - w = max(len(n) for n in result.endpoints) - for name, endpoint_id in sorted(result.endpoints.items()): - ui.console.print( - f" [white]{name:<{w}}[/white] " - f"{ui.endpoint_link(endpoint_id)}" + ui.resources_table( + ( + name, + handle.spec.kind.value, + ",".join(handle.spec.cpu or handle.spec.gpu or ["any"]), + result.endpoints.get(name, "registered"), ) + for name, handle in sorted(found.resources.items()) + ) ui.console.print() 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_deploy.py b/tests/test_apps/test_deploy.py index 5445c5b52..a29d3e405 100644 --- a/tests/test_apps/test_deploy.py +++ b/tests/test_apps/test_deploy.py @@ -191,7 +191,9 @@ def test_excluded_paths_are_logged_without_contents(self, tmp_path, caplog): with tarfile.open(package_project(tmp_path, {})) as tar: assert set(tar.getnames()) == {".runpodignore", "runpod_manifest.json"} for name in (".env", "private.key", "private.pem", "local.txt", ".aws/"): - assert f"excluded source path {name!r} from deployment artifact" in caplog.text + assert ( + f"excluded source path {name!r} from deployment artifact" in caplog.text + ) assert "sensitive-file-contents" not in caplog.text def test_vendored_env_included_under_env(self, tmp_path): @@ -280,7 +282,10 @@ async def test_deploy_app_full_flow(self, tmp_path): "uploadUrl": "https://upload", "objectKey": "key-1", } - api.finalize_artifact_upload.return_value = {"id": "build-1"} + api.finalize_artifact_upload.return_value = { + "id": "build-1", + "resourceSpec": {"tasks": {}}, + } api.deploy_build.return_value = {"id": "env-1"} with _stub_build(tmp_path): @@ -305,7 +310,10 @@ async def test_deploy_reuses_existing_app_and_env(self, tmp_path): "uploadUrl": "https://upload", "objectKey": "key-1", } - api.finalize_artifact_upload.return_value = {"id": "build-2"} + api.finalize_artifact_upload.return_value = { + "id": "build-2", + "resourceSpec": {"tasks": {}}, + } with _stub_build(tmp_path): await deploy_app(app, tmp_path, api=api) @@ -329,9 +337,18 @@ async def test_deploy_environment_is_used_for_nested_calls( "uploadUrl": "https://upload", "objectKey": "key-1", } - api.finalize_artifact_upload.return_value = {"id": "build-1"} + api.finalize_artifact_upload.return_value = { + "id": "build-1", + "resourceSpec": {"tasks": {}}, + } api.save_endpoint.return_value = {"id": "ep-1"} with _stub_build(tmp_path): + with pytest.raises(InvalidResourceError, match="FLASH_ENVIRONMENT"): + await deploy_app(app, tmp_path, env_name="prod", api=api) + api.get_app_by_name.assert_not_awaited() + api.deploy_build.assert_not_awaited() + api.save_endpoint.assert_not_awaited() + handle.spec.env = {} await deploy_app(app, tmp_path, env_name="prod", api=api) payload = api.save_endpoint.await_args.args[0] worker_env = { diff --git a/tests/test_apps/test_deployed_tasks.py b/tests/test_apps/test_deployed_tasks.py new file mode 100644 index 000000000..4fa5ce1ae --- /dev/null +++ b/tests/test_apps/test_deployed_tasks.py @@ -0,0 +1,471 @@ +"""production task identity, lifecycle, and deployment contract tests.""" + +import asyncio +from unittest.mock import AsyncMock, patch + +import pytest + +import runpod +from runpod.apps import App, Task +from runpod.apps.api import FlashTaskApiClient +from runpod.apps.app import _clear_registry +from runpod.apps.build import BuildResult +from runpod.apps.deploy import ( + _deployed_endpoint_input, + _prepare_task_resources, + build_manifest, + deploy_app, +) +from runpod.apps.errors import InvalidResourceError +from runpod.apps.targets import DeployedTaskTarget, PodTarget, _headers +from runpod.apps.tasks import DeployedTaskExecution, _pod_input +from runpod.apps.utils.network import runtime_headers +from runpod.apps.volume import GlobalVolume, NetworkVolume + + +@pytest.fixture(autouse=True) +def isolated_identity(monkeypatch): + _clear_registry() + for name in ( + "RUNPOD_RUNTIME_KIND", + "RUNPOD_BUILD_ID", + "RUNPOD_ENDPOINT_ID", + "RUNPOD_POD_ID", + "RUNPOD_AI_API_KEY", + "RUNPOD_API_KEY", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_DEV_APP", + "RUNPOD_DEV_SESSION", + ): + monkeypatch.delenv(name, raising=False) + yield + _clear_registry() + + +@pytest.fixture +def production(monkeypatch): + values = { + "RUNPOD_RUNTIME_KIND": "queue", + "RUNPOD_BUILD_ID": "build-1", + "FLASH_APP_NAME": "my-app", + "FLASH_ENVIRONMENT": "prod", + "FLASH_RESOURCE_NAME": "work", + "RUNPOD_ENDPOINT_ID": "ep-1", + "RUNPOD_POD_ID": "worker-pod", + "RUNPOD_AI_API_KEY": "endpoint-key", + "RUNPOD_API_KEY": "must-not-use-account-key", + } + for name, value in values.items(): + monkeypatch.setenv(name, value) + + +def test_endpoint_auth_emits_only_endpoint_source(production): + headers = _headers() + assert headers["Authorization"] == "Bearer endpoint-key" + assert headers["X-RunPod-Endpoint-Id"] == "ep-1" + assert "X-RunPod-Pod-Id" not in headers + + +def test_task_auth_emits_only_pod_source(production, monkeypatch): + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "task") + monkeypatch.setenv("RUNPOD_API_KEY", "child-key") + assert runtime_headers() == { + "Authorization": "Bearer child-key", + "X-RunPod-Pod-Id": "worker-pod", + } + + +@pytest.mark.parametrize("missing", ["RUNPOD_POD_ID", "RUNPOD_API_KEY"]) +def test_task_auth_never_uses_endpoint_or_global_key(production, monkeypatch, missing): + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "task") + monkeypatch.setattr(runpod, "api_key", "account-key") + monkeypatch.delenv(missing) + with pytest.raises(RuntimeError): + runtime_headers() + + +def test_endpoint_auth_does_not_fall_back_to_pod_key(production, monkeypatch): + monkeypatch.delenv("RUNPOD_AI_API_KEY") + monkeypatch.setattr(runpod, "api_key", "account-key") + with pytest.raises(RuntimeError, match="credential"): + runtime_headers() + + +@pytest.mark.parametrize( + "missing", + [ + "RUNPOD_AI_API_KEY", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_BUILD_ID", + ], +) +def test_missing_production_identity_never_falls_back(production, monkeypatch, missing): + monkeypatch.delenv(missing) + with pytest.raises(RuntimeError): + _headers() + + +def test_dev_auth_retains_account_key(monkeypatch): + monkeypatch.setenv("RUNPOD_DEV_APP", "my-app") + monkeypatch.setenv("RUNPOD_AI_API_KEY", "endpoint-key") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep-1") + monkeypatch.setenv("RUNPOD_API_KEY", "developer-key") + headers = _headers() + assert headers["Authorization"] == "Bearer developer-key" + assert "X-RunPod-Endpoint-Id" not in headers + + +@pytest.mark.parametrize("image", [None, "custom/image:tag"]) +def test_dev_task_payload_preserves_source_runtime(image): + app = App("my-app") + + @app.task(name="work", cpu="cpu3c-1-2", image=image) + def work(): + pass + + payload = _pod_input(work.spec, "execution-token", "work") + env = {entry["key"]: entry["value"] for entry in payload["env"]} + assert "RUNPOD_RUNTIME_KIND" not in env + assert "RUNPOD_BUILD_ID" not in env + assert "FLASH_APP_NAME" not in env + if image: + assert "runpod_sdk_runtime.task.runner" in payload["dockerArgs"] + + +async def test_graphql_create_and_terminate_use_exact_contract(production): + api = FlashTaskApiClient() + create_input = { + "flashAppName": "my-app", + "flashEnvironmentName": "prod", + "resourceName": "work", + "executionToken": "execution-token", + } + execute = AsyncMock( + side_effect=[ + { + "data": { + "flashTaskCreate": { + "podId": "child", + "flashBuildId": "build", + "resourceName": "work", + } + } + }, + {"data": {"flashTaskTerminate": True}}, + ] + ) + with patch("runpod.apps.api.run_graphql_query_async", execute): + assert (await api.create(create_input))["podId"] == "child" + await api.terminate("child", "execution-token") + create, terminate = execute.await_args_list + assert "FlashTaskCreateInput!" in create.args[0] + assert create.kwargs["variables"] == {"input": create_input} + assert create.kwargs["api_key"] == "endpoint-key" + assert create.kwargs["headers"]["X-RunPod-Endpoint-Id"] == "ep-1" + assert terminate.kwargs["variables"] == { + "input": {"podId": "child", "executionToken": "execution-token"} + } + + +async def test_explicit_local_task_uses_account_named_api(monkeypatch): + monkeypatch.setenv("RUNPOD_API_KEY", "local-key") + execute = AsyncMock(return_value={"data": {"flashTaskCreate": {"podId": "child"}}}) + with patch("runpod.apps.api.run_graphql_query_async", execute): + await FlashTaskApiClient().create({"resourceName": "work"}) + assert execute.await_args.kwargs["api_key"] == "local-key" + assert execute.await_args.kwargs["headers"] is None + + +async def test_task_target_selection_is_production_only(production, monkeypatch): + app = App("my-app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(value): + return value + + target = await app._resolve(work.spec) + assert isinstance(target, DeployedTaskTarget) + assert target.env_name == "prod" + assert target.build_payload(work._fn, work.spec, (1,), {"other": 2}) == { + "args": [1], + "kwargs": {"other": 2}, + } + for name in ("RUNPOD_BUILD_ID", "FLASH_APP_NAME", "FLASH_RESOURCE_NAME"): + monkeypatch.delenv(name) + assert isinstance(await app._resolve(work.spec), PodTarget) + monkeypatch.setenv("RUNPOD_DEV_SESSION", "session") + assert type(await app._resolve(work.spec)) is PodTarget + + +async def test_invalid_production_task_identity_cannot_use_source_target( + production, monkeypatch +): + app = App("my-app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(): + pass + + monkeypatch.delenv("FLASH_APP_NAME") + with pytest.raises(RuntimeError, match="identity"): + await app._resolve(work.spec) + + +@pytest.mark.parametrize("sync", [False, True]) +def test_task_self_remote_launches_child(production, sync): + app = App("my-app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(value): + raise AssertionError("task body must not run inline") + + target = DeployedTaskTarget("my-app", "prod", "work") + target.invoke = AsyncMock(return_value="child-result") + with patch.object(app, "_resolve", AsyncMock(return_value=target)): + result = work.remote(1) if sync else asyncio.run(work.remote.aio(1)) + assert result == "child-result" + target.invoke.assert_awaited_once_with({"args": [1], "kwargs": {}}) + + +def test_production_queue_self_optimization_checks_app(production): + other = App("other-app") + + @other.queue(name="work", cpu="cpu3c-1-2") + def work(): + pass + + assert not work._is_current_worker() + assert work._app.name != "my-app" + + +async def test_task_stub_remote_and_spawn_are_json_only(): + stub = Task(app="my-app", env="prod", name="work") + assert runpod.Task is Task + target = stub._target() + target.invoke = AsyncMock(return_value=3) + target.submit = AsyncMock(return_value="job") + with patch.object(stub, "_target", return_value=target): + assert await stub.remote.aio(1, value=2) == 3 + assert await stub.spawn.aio(1, value=2) == "job" + target.invoke.assert_awaited_once_with({"args": [1], "kwargs": {"value": 2}}) + target.submit.assert_awaited_once_with({"args": [1], "kwargs": {"value": 2}}) + + +@pytest.mark.parametrize( + "payload", + [ + {"args": [], "kwargs": {}, "function_code": "bad"}, + {"args": "pickle", "kwargs": {}}, + {"args": [object()], "kwargs": {}}, + {"args": [float("nan")], "kwargs": {}}, + {"args": [], "kwargs": {1: "bad"}}, + ], +) +async def test_deployed_payload_rejects_overrides_and_non_json(payload): + target = DeployedTaskTarget("my-app", "prod", "work") + with pytest.raises((ValueError, TypeError)): + await target.submit(payload) + + +async def test_cancellation_retains_pending_create_and_cleans_success(): + entered = asyncio.Event() + finish = asyncio.Event() + + async def create(task_input): + entered.set() + await finish.wait() + return {"podId": "child"} + + api = AsyncMock() + api.create.side_effect = create + execution = DeployedTaskExecution("my-app", "prod", "work", api) + target = DeployedTaskTarget("my-app", "prod", "work") + with patch.object(target, "_execution", return_value=execution): + invocation = asyncio.create_task(target.submit({"args": [], "kwargs": {}})) + await entered.wait() + invocation.cancel() + await asyncio.sleep(0) + assert not invocation.done() + finish.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(invocation, timeout=5) + api.create.assert_awaited_once() + api.terminate.assert_awaited_once_with("child", execution.token) + assert execution.pod_id is None + + +async def test_create_failure_is_not_retried_or_source_fallback(): + api = AsyncMock() + api.create.side_effect = RuntimeError("approval missing") + execution = DeployedTaskExecution("my-app", "prod", "work", api) + target = DeployedTaskTarget("my-app", "prod", "work") + with patch.object(target, "_execution", return_value=execution): + with pytest.raises(RuntimeError, match="approval missing"): + await target.invoke({"args": [], "kwargs": {}}) + api.create.assert_awaited_once() + api.terminate.assert_not_awaited() + + +async def test_execution_creation_contains_only_named_identity_and_timeout(): + api = AsyncMock() + api.create.return_value = {"podId": "child"} + execution = DeployedTaskExecution("my-app", "prod", "work", api, timeout=12.5) + await execution.start() + api.create.assert_awaited_once_with( + { + "flashAppName": "my-app", + "flashEnvironmentName": "prod", + "resourceName": "work", + "executionToken": execution.token, + "executionTimeoutSeconds": 12.5, + } + ) + assert execution._headers == {"Authorization": f"Bearer {execution.token}"} + await execution.terminate() + + +def test_manifest_pins_task_image_and_full_storage_references(tmp_path, monkeypatch): + monkeypatch.setenv("RUNPOD_RUNTIME_TAG", "release") + app = App("my-app") + + @app.task( + name="work", + gpu="4090", + mounts={ + "/data": NetworkVolume("data", size=75, datacenter="EU-RO-1", create=False), + "/global": GlobalVolume("shared", create=True), + }, + ) + def work(): + pass + + resource = build_manifest(app, tmp_path, python_version="3.12")["resources"][0] + assert resource["imageName"] == "runpod/task-gpu:py3.12-release" + assert resource["mounts"] == [ + { + "kind": "network", + "reference": "data", + "path": "/data", + "create": False, + "size": 75, + "datacenter": "EU-RO-1", + }, + {"kind": "global", "reference": "shared", "path": "/global", "create": True}, + ] + assert resource["module"] + assert resource["qualname"].endswith("work") + + +def test_deployed_endpoint_uses_runtime_identity_without_account_key(monkeypatch): + monkeypatch.setattr(runpod, "api_key", "developer-key") + app = App("my-app") + + @app.queue(name="work", cpu="cpu3c-1-2") + def work(): + pass + + payload = _deployed_endpoint_input(app, work.spec, "env", "build", "3.12") + env = {e["key"]: e["value"] for e in payload["template"]["env"]} + assert env["FLASH_APP_NAME"] == "my-app" + assert "RUNPOD_API_KEY" not in env + assert env["RUNPOD_BUILD_ID"] == "build" + assert env["FLASH_RESOURCE_NAME"] == "work" + assert env["FLASH_ENVIRONMENT"] == "default" + work.spec.env = {"RUNPOD_AI_API_KEY": "override"} + with pytest.raises(InvalidResourceError, match="reserved"): + _deployed_endpoint_input(app, work.spec, "env", "build", "3.12") + + +async def test_task_storage_and_registry_are_resolved_with_deployer_client(tmp_path): + app = App("my-app") + + @app.task( + name="work", + cpu="cpu3c-1-2", + image="private/image:tag", + registry_auth="private", + mounts={"/data": NetworkVolume("data", size=75, create=True)}, + ) + def work(): + pass + + manifest = build_manifest(app, tmp_path) + api = AsyncMock() + api.list_registry_auths.return_value = [{"id": "registry-id", "name": "private"}] + bindings = [ + { + "id": "volume-id", + "kind": "network", + "reference": "data", + "path": "/data", + "dataCenterId": "EU-RO-1", + } + ] + with patch( + "runpod.apps.volume.VolumeResolver.resolve_mounts", + AsyncMock(return_value=bindings), + ): + await _prepare_task_resources(app, manifest, api) + resource = manifest["resources"][0] + assert resource["mounts"][0]["id"] == "volume-id" + assert resource["mounts"][0]["reference"] == "data" + assert resource["mounts"][0]["size"] == 75 + assert resource["registryAuth"] == "registry-id" + assert resource["customImage"] is True + api.list_registry_auths.assert_awaited_once() + + +@pytest.mark.parametrize( + "definitions,approved", + [ + (None, False), + ({}, False), + ({"tasks": {}}, False), + ({"tasks": {"work": None}}, False), + ({"tasks": {"work": {}}}, False), + ({"tasks": {"other": {"name": "other"}}}, False), + ({"tasks": {"work": {"name": "work"}, "extra": {"name": "extra"}}}, False), + ({"tasks": {"work": {"name": "work"}}}, True), + ], +) +async def test_task_only_deployment_requires_definitions_without_launching( + tmp_path, definitions, approved +): + app = App("my-app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(): + pass + + api = AsyncMock() + api.get_app_by_name.return_value = { + "id": "app", + "flashEnvironments": [{"id": "env", "name": "default"}], + } + api.prepare_artifact_upload.return_value = { + "uploadUrl": "upload", + "objectKey": "key", + } + api.finalize_artifact_upload.return_value = {"id": "build"} + api.finalize_artifact_upload.return_value["resourceSpec"] = definitions + env_dir = tmp_path / "vendor" + env_dir.mkdir() + with patch( + "runpod.apps.deploy.build_environment", + return_value=BuildResult(env_dir=env_dir, requirements=[]), + ): + if approved: + result = await deploy_app(app, tmp_path, api=api) + assert result.resources == ["work"] + assert result.endpoints == {} + api.deploy_build.assert_awaited_once_with("env", "build") + else: + with pytest.raises(RuntimeError, match="acknowledge"): + await deploy_app(app, tmp_path, api=api) + api.deploy_build.assert_not_awaited() + api.save_endpoint.assert_not_awaited() + api.deploy_task_pod.assert_not_awaited() diff --git a/tests/test_apps/test_deployment_validation.py b/tests/test_apps/test_deployment_validation.py new file mode 100644 index 000000000..7853c8523 --- /dev/null +++ b/tests/test_apps/test_deployment_validation.py @@ -0,0 +1,345 @@ +"""deployment policy validation and shared resource resolution.""" + +import json +import tarfile +from unittest.mock import AsyncMock, patch + +import pytest + +from runpod.apps import App +from runpod.apps.api import AppsApiClient +from runpod.apps.app import _clear_registry +from runpod.apps.build import BuildResult +from runpod.apps.deploy import ( + _deployed_endpoint_input, + build_artifact, + build_manifest, + deploy_app, +) +from runpod.apps.errors import InvalidResourceError +from runpod.apps.placement import solve_placement +from runpod.apps.tasks import DeployedTaskExecution, _validate_timeout +from runpod.apps.volume import GlobalVolume, NetworkVolume + + +@pytest.fixture(autouse=True) +def isolated_app_registry(): + _clear_registry() + yield + _clear_registry() + + +@pytest.fixture +def deployment_api(): + api = AsyncMock() + api.get_app_by_name.return_value = { + "id": "app", + "flashEnvironments": [{"id": "env", "name": "default"}], + } + api.prepare_artifact_upload.return_value = { + "uploadUrl": "upload", + "objectKey": "key", + } + api.finalize_artifact_upload.return_value = {"id": "build"} + api.save_endpoint.return_value = {"id": "endpoint"} + return api + + +@pytest.mark.parametrize( + "key", + [ + "RUNPOD_API_KEY", + "runpod_custom", + "FLASH_APP_NAME", + " flash_custom ", + "PYTHONPATH", + "pythonhome", + "LD_PRELOAD", + ], +) +def test_task_manifest_rejects_reserved_env_keys(tmp_path, key): + app = App("policy-app") + + @app.task(name="work", cpu="cpu3c-1-2", env={key: "test-value"}) + def work(): + pass + + with pytest.raises(InvalidResourceError) as exc: + build_manifest(app, tmp_path) + assert key in str(exc.value) + assert "test-value" not in str(exc.value) + + +@pytest.mark.parametrize( + "path", + [ + "/", + "/root", + "/root/.runpod", + "/root/.runpod/cache", + "/runpod-volume", + "/runpod-volume/data", + "/proc", + "/proc/data", + "/sys", + "/sys/data", + "/dev", + "/dev/data", + ], +) +def test_task_manifest_rejects_reserved_mount_overlap(tmp_path, path): + app = App("policy-app") + + @app.task(name="work", cpu="cpu3c-1-2", mounts={path: NetworkVolume("data")}) + def work(): + pass + + with pytest.raises(InvalidResourceError, match="reserved runtime path"): + build_manifest(app, tmp_path) + + +@pytest.mark.parametrize("path", ["/data", "/root/.runpod-data", "/proc-data"]) +def test_task_manifest_accepts_non_reserved_env_and_paths(tmp_path, path): + app = App("policy-app") + + @app.task( + name="work", + cpu="cpu3c-1-2", + mounts={path: NetworkVolume("data")}, + env={"APP_TOKEN": "test-value", "PYTHONPATH_EXTRA": "allowed"}, + ) + def work(): + pass + + resource = build_manifest(app, tmp_path)["resources"][0] + assert resource["mounts"][0]["path"] == path + assert resource["env"] == work.spec.env + + +@pytest.mark.parametrize( + "options", + [ + {"env": {"RUNPOD_API_KEY": "test-value"}}, + {"mounts": {"/root": NetworkVolume("data")}}, + ], +) +async def test_invalid_tasks_fail_before_vendoring_or_remote_calls(tmp_path, options): + app = App("policy-app") + + @app.task(name="work", cpu="cpu3c-1-2", **options) + def work(): + pass + + api = AsyncMock() + with patch("runpod.apps.deploy.build_environment") as vendor: + with pytest.raises(InvalidResourceError): + build_artifact(app, tmp_path) + with pytest.raises(InvalidResourceError): + await deploy_app(app, tmp_path, api=api) + vendor.assert_not_called() + assert not api.method_calls + + +@pytest.mark.parametrize("kind", ["queue", "api"]) +@pytest.mark.parametrize( + "key", + [ + "RUNPOD_API_KEY", + "RUNPOD_AI_API_KEY", + "RUNPOD_ENDPOINT_ID", + "RUNPOD_POD_ID", + "RUNPOD_TASK_TOKEN", + "RUNPOD_TASK_PORT", + "FLASH_BUILD_ARTIFACT_PATH", + "RUNPOD_RUNTIME_KIND", + "RUNPOD_DEV_APP", + "RUNPOD_DEV_RESOURCE", + "RUNPOD_DEV_SESSION", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_BUILD_ID", + ], +) +async def test_later_endpoint_reserved_env_fails_before_any_remote_change( + tmp_path, deployment_api, kind, key +): + app = App("policy-app") + + @app.queue(name="alpha", cpu="cpu3c-1-2") + def alpha(): + pass + + @getattr(app, kind)(name="zulu", cpu="cpu3c-1-2", env={key: "test-value"}) + def zulu(): + pass + + vendor_dir = tmp_path / "vendor" + vendor_dir.mkdir() + with patch( + "runpod.apps.deploy.build_environment", + return_value=BuildResult(env_dir=vendor_dir), + ) as vendor: + with pytest.raises(InvalidResourceError, match=key): + await deploy_app(app, tmp_path, api=deployment_api) + vendor.assert_not_called() + deployment_api.upload_tarball.assert_not_awaited() + deployment_api.deploy_build.assert_not_awaited() + deployment_api.save_endpoint.assert_not_awaited() + deployment_api.delete_endpoint.assert_not_awaited() + assert not deployment_api.method_calls + with pytest.raises(InvalidResourceError, match=key): + _deployed_endpoint_input(app, zulu.spec, "env", "build", "3.12") + + +def test_endpoint_reserved_env_error_identifies_keys_without_values(): + app = App("policy-app") + + @app.queue( + name="worker", + cpu="cpu3c-1-2", + env={"RUNPOD_API_KEY": "test-value", "RUNPOD_AI_API_KEY": "other-value"}, + ) + def worker(): + pass + + with pytest.raises(InvalidResourceError) as exc: + _deployed_endpoint_input(app, worker.spec, "env", "build", "3.12") + assert "RUNPOD_API_KEY" in str(exc.value) + assert "RUNPOD_AI_API_KEY" in str(exc.value) + assert "test-value" not in str(exc.value) + assert "other-value" not in str(exc.value) + + +@pytest.mark.parametrize("timeout", [0, -1, True, "1", float("inf"), float("nan")]) +def test_deployed_task_rejects_invalid_timeout_before_create(timeout): + api = AsyncMock() + with pytest.raises(ValueError, match="timeout"): + DeployedTaskExecution("app", "env", "work", api, timeout=timeout) + api.create.assert_not_called() + + +@pytest.mark.parametrize("timeout", [None, 0.5, 10]) +async def test_deployed_task_accepts_positive_or_omitted_timeout(timeout): + api = AsyncMock() + api.create.return_value = {"podId": "pod"} + execution = DeployedTaskExecution("app", "env", "work", api, timeout=timeout) + await execution.start() + request = api.create.await_args.args[0] + if timeout is None: + assert "executionTimeoutSeconds" not in request + else: + assert request["executionTimeoutSeconds"] == timeout + + +def test_transport_still_accepts_zero_timeout(): + _validate_timeout(0) + + +@pytest.mark.parametrize("kind", [None, "queue", "task"]) +async def test_finalize_selects_task_policy_only_for_task_manifests(kind): + api = AppsApiClient(api_key="test-key") + manifest = {"resources": [] if kind is None else [{"name": "work", "kind": kind}]} + response = {"id": "build"} + with patch.object( + api, + "_execute", + AsyncMock(return_value={"finalizeFlashArtifactUpload": response}), + ) as execute: + assert await api.finalize_artifact_upload("app", "key", manifest) == response + query, variables = execute.await_args.args + assert ("resourceSpec" in query) is (kind == "task") + assert variables["input"]["manifest"] == manifest + + +async def test_taskless_deployment_accepts_finalize_without_task_policy( + tmp_path, deployment_api +): + app = App("policy-app") + + @app.queue(name="worker", cpu="cpu3c-1-2") + def worker(): + pass + + vendor_dir = tmp_path / "vendor" + vendor_dir.mkdir() + with patch( + "runpod.apps.deploy.build_environment", + return_value=BuildResult(env_dir=vendor_dir, excluded=["torch"]), + ): + result = await deploy_app(app, tmp_path, api=deployment_api) + assert result.endpoints == {"worker": "endpoint"} + deployment_api.deploy_build.assert_awaited_once_with("env", "build") + manifest = deployment_api.finalize_artifact_upload.await_args.args[2] + assert manifest["excludedPackages"] == ["torch"] + + +def test_build_artifact_preserves_excluded_packages(tmp_path): + app = App("policy-app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(): + pass + + vendor_dir = tmp_path / "vendor" + vendor_dir.mkdir() + with patch( + "runpod.apps.deploy.build_environment", + return_value=BuildResult(env_dir=vendor_dir, excluded=["torch"]), + ): + artifact = build_artifact(app, tmp_path) + with tarfile.open(artifact) as tar: + manifest = json.load(tar.extractfile("runpod_manifest.json")) + assert manifest["excludedPackages"] == ["torch"] + + +@pytest.mark.parametrize("volume_type", [GlobalVolume, NetworkVolume]) +async def test_deploy_resolves_shared_task_endpoint_volume_once( + tmp_path, deployment_api, volume_type +): + app = App("policy-app") + volume = volume_type("shared") + + @app.task(name="work", gpu="4090", mounts={"/data": volume}) + def work(): + pass + + @app.queue(name="worker", gpu="4090", mounts={"/runpod-volume": volume}) + def worker(): + pass + + deployment_api.list_global_volumes.return_value = [ + {"id": "volume", "name": "shared"} + ] + deployment_api.list_network_volumes.return_value = [ + {"id": "volume", "name": "shared", "dataCenter": "EU-RO-1"} + ] + deployment_api.gpu_stock_status.return_value = "High" + deployment_api.finalize_artifact_upload.return_value["resourceSpec"] = { + "tasks": {"work": {"imageName": "approved"}} + } + vendor_dir = tmp_path / "vendor" + vendor_dir.mkdir() + with ( + patch( + "runpod.apps.deploy.build_environment", + return_value=BuildResult(env_dir=vendor_dir), + ), + patch("runpod.apps.placement.solve_placement", wraps=solve_placement) as solve, + ): + await deploy_app(app, tmp_path, api=deployment_api) + manifest = deployment_api.finalize_artifact_upload.await_args.args[2] + task = next(r for r in manifest["resources"] if r["kind"] == "task") + assert task["mounts"][0]["id"] == "volume" + endpoint = deployment_api.save_endpoint.await_args.args[0] + if volume_type is GlobalVolume: + deployment_api.list_global_volumes.assert_awaited_once() + solve.assert_not_called() + assert endpoint["volumes"] == [ + {"volumeId": "volume", "volumeType": "OBJECT_STORE_VOLUME"} + ] + else: + deployment_api.list_network_volumes.assert_awaited_once() + solve.assert_called_once() + assert endpoint["networkVolumeIds"] == [{"networkVolumeId": "volume"}] + assert endpoint["locations"] == "EU-RO-1" diff --git a/tests/test_apps/test_dispatch.py b/tests/test_apps/test_dispatch.py index 6a2543a83..b2967a418 100644 --- a/tests/test_apps/test_dispatch.py +++ b/tests/test_apps/test_dispatch.py @@ -207,6 +207,9 @@ def t(x): def test_worker_executes_own_body(self, monkeypatch): monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep-1") monkeypatch.setenv("FLASH_RESOURCE_NAME", "q") + monkeypatch.setenv("FLASH_APP_NAME", "my-app") + monkeypatch.setenv("FLASH_ENVIRONMENT", "production") + monkeypatch.setenv("RUNPOD_BUILD_ID", "build-1") app = App("my-app") @app.queue(name="q") @@ -251,6 +254,9 @@ async def main(): def test_worker_calls_sibling_via_sentinel(self, monkeypatch): monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep-1") monkeypatch.setenv("FLASH_RESOURCE_NAME", "other") + monkeypatch.setenv("FLASH_APP_NAME", "my-app") + monkeypatch.setenv("FLASH_ENVIRONMENT", "production") + monkeypatch.setenv("RUNPOD_BUILD_ID", "build-1") app = App("my-app") @app.queue(name="q") @@ -546,19 +552,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_runtime_call_validation.py b/tests/test_apps/test_runtime_call_validation.py new file mode 100644 index 000000000..782618c3d --- /dev/null +++ b/tests/test_apps/test_runtime_call_validation.py @@ -0,0 +1,172 @@ +"""runtime scope enforcement and task argument compatibility.""" + +from dataclasses import dataclass +from datetime import datetime +from unittest.mock import AsyncMock, patch + +import pytest + +import runpod +from runpod.apps import App +from runpod.apps.api import FlashTaskApiClient +from runpod.apps.app import _clear_registry +from runpod.apps.targets import DeployedTaskTarget, PodTarget, SentinelTarget + + +@pytest.fixture(autouse=True) +def isolated_runtime(monkeypatch): + _clear_registry() + for key in ( + "RUNPOD_RUNTIME_KIND", + "RUNPOD_BUILD_ID", + "RUNPOD_ENDPOINT_ID", + "RUNPOD_POD_ID", + "RUNPOD_AI_API_KEY", + "RUNPOD_API_KEY", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_DEV_APP", + "RUNPOD_DEV_SESSION", + ): + monkeypatch.delenv(key, raising=False) + yield + _clear_registry() + + +@pytest.fixture(params=["queue", "task"]) +def production(request, monkeypatch): + for key, value in { + "RUNPOD_RUNTIME_KIND": request.param, + "RUNPOD_BUILD_ID": "build", + "RUNPOD_ENDPOINT_ID": "endpoint", + "RUNPOD_POD_ID": "pod", + "RUNPOD_AI_API_KEY": "endpoint-key", + "RUNPOD_API_KEY": "pod-key", + "FLASH_APP_NAME": "my-app", + "FLASH_ENVIRONMENT": "prod", + "FLASH_RESOURCE_NAME": "caller", + }.items(): + monkeypatch.setenv(key, value) + monkeypatch.setattr(runpod, "api_key", "account-key") + return request.param + + +@pytest.mark.parametrize("kind", ["queue", "queue-spawn", "api", "task", "task-spawn"]) +@pytest.mark.parametrize("app,env", [("other-app", "prod"), ("my-app", "staging")]) +async def test_cross_scope_stubs_fail_before_network_request( + production, kind, app, env +): + with ( + patch("runpod.apps.targets._request_json", AsyncMock()) as sentinel, + patch("runpod.apps.api.run_graphql_query_async", AsyncMock()) as graphql, + ): + with pytest.raises(ValueError, match="same app and environment"): + if kind == "api": + await runpod.Api(app=app, env=env, name="work").post.aio("/work", {}) + elif kind.startswith("queue"): + stub = runpod.Queue(app=app, env=env, name="work") + await (stub.spawn if kind.endswith("spawn") else stub.remote).aio() + else: + stub = runpod.Task(app=app, env=env, name="work") + await (stub.spawn if kind.endswith("spawn") else stub.remote).aio() + sentinel.assert_not_awaited() + graphql.assert_not_awaited() + + +async def test_same_scope_targets_use_runtime_credentials(production): + key = "endpoint-key" if production == "queue" else "pod-key" + headers = SentinelTarget("my-app", "prod", "work")._sentinel_headers() + assert headers["Authorization"] == f"Bearer {key}" + execute = AsyncMock(return_value={"data": {"flashTaskCreate": {"podId": "child"}}}) + with patch("runpod.apps.api.run_graphql_query_async", execute): + await FlashTaskApiClient().create( + { + "flashAppName": "my-app", + "flashEnvironmentName": "prod", + "resourceName": "work", + } + ) + assert execute.await_args.kwargs["api_key"] == key + + +@pytest.mark.parametrize("dev", [False, True]) +async def test_account_clients_can_target_other_apps_and_environments(monkeypatch, dev): + monkeypatch.setenv("RUNPOD_API_KEY", "account-key") + if dev: + monkeypatch.setenv("RUNPOD_DEV_SESSION", "dev-session") + monkeypatch.setenv("RUNPOD_DEV_APP", "my-app") + headers = SentinelTarget("other-app", "staging", "work")._sentinel_headers() + assert headers["Authorization"] == "Bearer account-key" + assert "X-RunPod-Endpoint-Id" not in headers + assert "X-RunPod-Pod-Id" not in headers + execute = AsyncMock(return_value={"data": {"flashTaskCreate": {"podId": "child"}}}) + with patch("runpod.apps.api.run_graphql_query_async", execute): + await FlashTaskApiClient().create( + { + "flashAppName": "other-app", + "flashEnvironmentName": "staging", + "resourceName": "work", + } + ) + assert execute.await_args.kwargs["api_key"] == "account-key" + assert execute.await_args.kwargs["headers"] is None + + +@dataclass +class Argument: + value: int + + +@pytest.mark.parametrize( + "value", [b"payload", Argument(1), datetime(2026, 1, 1), {1, 2}, float("nan")] +) +@pytest.mark.parametrize("method", ["invoke", "submit"]) +async def test_deployed_task_reports_consistent_json_error_before_creation( + value, method +): + target = DeployedTaskTarget("app", "env", "work") + with patch("runpod.apps.api.run_graphql_query_async", AsyncMock()) as execute: + with pytest.raises( + ValueError, match="^tasks accept only JSON args and kwargs$" + ): + await getattr(target, method)({"args": [value], "kwargs": {}}) + execute.assert_not_awaited() + + +@pytest.mark.parametrize("dev", [False, True]) +@pytest.mark.parametrize( + "value", [b"private-payload", Argument(1), datetime(2026, 1, 1)] +) +async def test_source_tasks_reject_non_json_arguments(monkeypatch, dev, value): + if dev: + monkeypatch.setenv("RUNPOD_DEV_SESSION", "dev-session") + app = App("app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(value): + return value + + target = await app._resolve(work.spec) + assert isinstance(target, PodTarget) + with pytest.raises( + ValueError, match="^tasks accept only JSON args and kwargs$" + ) as exc: + target.build_payload(work._fn, work.spec, (value,), {}) + assert "private-payload" not in str(exc.value) + assert work.local(value) == value + + +def test_source_tasks_use_json_arguments_and_results(): + app = App("app") + + @app.task(name="work", cpu="cpu3c-1-2") + def work(value): + return value + + value = {"items": [1, "two", True, None]} + payload = PodTarget(work.spec, work._fn).build_payload( + work._fn, work.spec, (value,), {} + ) + assert payload["args"] == [value] + assert payload["serialization_format"] == "json" diff --git a/tests/test_apps/test_tasks.py b/tests/test_apps/test_tasks.py index 78b761152..f20934bc3 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,8 @@ 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["serialization_format"] == "json" assert payload["dependencies"] == ["numpy"] def test_task_remote_runs_past_one_hour_and_cleans_up(self): @@ -159,7 +154,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_apps/test_transport_parity.py b/tests/test_apps/test_transport_parity.py new file mode 100644 index 000000000..55fb99b24 --- /dev/null +++ b/tests/test_apps/test_transport_parity.py @@ -0,0 +1,170 @@ +"""json contracts across development and deployed invocation targets.""" + +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from runpod.apps import App, Context +from runpod.apps.app import _clear_registry +from runpod.apps.spec import ResourceKind, ResourceSpec +from runpod.apps.targets import ( + DeployedTaskTarget, + LiveTarget, + PodTarget, + SentinelTarget, + _request_json, +) +from runpod.apps.tasks import unwrap_task_response + + +@pytest.fixture(autouse=True) +def isolated_apps(): + _clear_registry() + yield + _clear_registry() + + +@pytest.mark.parametrize( + "target", [LiveTarget("dev"), SentinelTarget("app", "env", "work")] +) +@pytest.mark.parametrize( + "value", [b"bytes", datetime(2026, 1, 1), float("nan"), float("inf")] +) +def test_queue_arguments_obey_same_json_contract(target, value): + def work(value): + return value + + spec = ResourceSpec(kind=ResourceKind.QUEUE, name="work") + with pytest.raises(TypeError, match="json-serializable"): + target.build_payload(work, spec, (value,), {}) + + +@pytest.mark.parametrize( + "value", [b"bytes", datetime(2026, 1, 1), float("nan"), float("inf")] +) +def test_task_arguments_obey_same_json_contract(value): + def work(value): + return value + + spec = ResourceSpec(kind=ResourceKind.TASK, name="work") + with pytest.raises(ValueError, match="only JSON args and kwargs"): + PodTarget(spec, work).build_payload(work, spec, (), {"value": value}) + with pytest.raises(ValueError, match="only JSON args and kwargs"): + DeployedTaskTarget("app", "env", "work")._execution( + {"args": [], "kwargs": {"value": value}} + ) + + +def test_dev_task_json_arguments_match_deployed_payload(): + def work(value, flag): + return value + + spec = ResourceSpec(kind=ResourceKind.TASK, name="work") + args, kwargs = ({"values": [1, None]},), {"flag": True} + dev = PodTarget(spec, work).build_payload(work, spec, args, kwargs) + deployed = DeployedTaskTarget("app", "env", "work").build_payload( + work, spec, args, kwargs + ) + assert dev["serialization_format"] == "json" + assert {key: dev[key] for key in ("args", "kwargs")} == deployed + + +@pytest.mark.parametrize( + "target", [LiveTarget("dev"), SentinelTarget("app", "env", "work")] +) +@pytest.mark.parametrize("value", [b"bytes", float("nan"), float("inf")]) +async def test_api_body_validation_precedes_requests_and_source_sync(target, value): + with ( + patch("runpod.apps.targets.aiohttp.ClientSession") as session, + patch.object(LiveTarget, "_sync_source", AsyncMock()) as sync, + ): + with pytest.raises( + ValueError, match="api request body must be JSON-serializable" + ): + await target.request("POST", "/work", {"value": value}) + session.assert_not_called() + sync.assert_not_awaited() + + +async def test_raw_queue_requests_reject_non_json_options_before_network(): + with patch("runpod.apps.targets.aiohttp.ClientSession") as session: + with pytest.raises(ValueError, match="remote request body"): + await _request_json( + "POST", + "https://invalid", + {}, + 30, + payload={"input": {}, "policy": {"ttl": float("nan")}}, + ) + session.assert_not_called() + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_queue_and_task_results_reject_non_finite_json(value): + for target in (LiveTarget("dev"), SentinelTarget("app", "env", "work")): + with pytest.raises(ValueError, match="results must be JSON-serializable"): + target.unwrap({"status": "COMPLETED", "output": {"value": value}}) + with pytest.raises(ValueError, match="results must be JSON-serializable"): + LiveTarget("dev")._unwrap_response({"success": True, "json_result": value}) + with pytest.raises(ValueError, match="results must be JSON-serializable"): + unwrap_task_response({"success": True, "json_result": value}) + + +@pytest.mark.parametrize( + "text,expected", [(' {"value": 1} ', {"value": 1}), ("plain text", "plain text")] +) +@pytest.mark.parametrize( + "target", [LiveTarget("dev"), SentinelTarget("app", "env", "work")] +) +async def test_api_json_and_text_responses_share_transport(target, text, expected): + with ( + patch.object(target._client, "_headers", return_value={}), + patch("runpod.apps.targets.aiohttp.ClientSession") as factory, + ): + session = factory.return_value.__aenter__.return_value + session.request = MagicMock() + response = session.request.return_value.__aenter__.return_value + response.status = 200 + response.text = AsyncMock(return_value=text) + response.raise_for_status = lambda: None + assert await target.request("POST", "/work", {"value": 1}) == expected + assert session.request.call_args.kwargs["json"] == {"value": 1} + + +@pytest.mark.parametrize("async_call", [False, True]) +async def test_worker_self_calls_apply_json_conversion(async_call): + app = App("app") + + @app.queue(name="work") + def work(value): + assert isinstance(value, list) + return (value,) + + with ( + patch("runpod.apps.handles.current_context", return_value=Context.WORKER), + patch.object(work, "_is_current_worker", return_value=True), + ): + result = await work.remote.aio((1, 2)) if async_call else work.remote((1, 2)) + assert result == [[1, 2]] + with pytest.raises(ValueError, match="JSON-serializable"): + if async_call: + await work.remote.aio(b"bytes") + else: + work.remote(b"bytes") + + +async def test_worker_self_stream_rejects_non_json_chunks(): + app = App("app") + + @app.queue(name="work") + def work(): + yield b"bytes" + + with ( + patch("runpod.apps.handles.current_context", return_value=Context.WORKER), + patch.object(work, "_is_current_worker", return_value=True), + ): + with pytest.raises(ValueError, match="JSON-serializable"): + async for _ in work.stream.aio(): + pass diff --git a/tests/test_cli/test_deploy_command.py b/tests/test_cli/test_deploy_command.py index 59ea111ee..14460c90e 100644 --- a/tests/test_cli/test_deploy_command.py +++ b/tests/test_cli/test_deploy_command.py @@ -63,6 +63,34 @@ def _result(name="demo", endpoints=None): class TestDeploy: + + def test_task_only_deployment_shows_registered_per_call( + self, tmp_path, monkeypatch + ): + source = """ +import runpod +app = runpod.App("tasks") +@app.task(name="batch", cpu="cpu3c-1-2") +def batch(): + return 1 +""" + (tmp_path / "main.py").write_text(source) + monkeypatch.chdir(tmp_path) + deployed = DeployResult( + app_name="tasks", + build_id="build", + environment_id="env", + resources=["batch"], + endpoints={}, + ) + with patch("runpod.apps.deploy.deploy_app", AsyncMock(return_value=deployed)): + result = _runner().invoke(cli, ["flash", "deploy"]) + assert result.exit_code == 0, result.output + assert "batch" in result.output + assert "task" in result.output + assert "cpu3c-1-2" in result.output + assert "registered / per-call" in result.output + def test_deploys_discovered_app(self, tmp_path, monkeypatch): (tmp_path / "main.py").write_text(APP_SOURCE) monkeypatch.chdir(tmp_path) diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 000000000..2fd6fe500 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,231 @@ +"""json normalization and response safety shared with worker runtimes.""" + +import json +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from runpod._serialization import normalize_json, validate_json +from runpod.apps.errors import RemoteExecutionError +from runpod.apps.protocol import FunctionRequest, FunctionResponse +from runpod.apps.targets import LiveTarget +from runpod.apps.tasks import unwrap_task_response +from runpod.serverless.modules import rp_fastapi, rp_http, rp_job + + +@pytest.mark.parametrize( + "value,expected", + [ + ((1, 2), [1, 2]), + ({1: (None, True)}, {"1": [None, True]}), + ({None: False}, {"null": False}), + (None, None), + ], +) +def test_shared_json_conversion_matches_wire(value, expected): + validate_json(value) + assert normalize_json(value) == expected + + +@pytest.mark.parametrize( + "value", + [b"bytes", datetime(2026, 1, 1), object(), float("nan"), float("inf"), {1, 2}], +) +def test_shared_validation_rejects_non_json_values(value): + with pytest.raises(TypeError, match="json-serializable"): + validate_json(value) + with pytest.raises(TypeError, match="json-serializable"): + normalize_json(value) + + +def test_shared_validation_rejects_cycles(): + value = [] + value.append(value) + with pytest.raises(TypeError, match="(?i)circular reference"): + validate_json(value) + + +def test_function_request_defaults_to_json_and_normalizes_values(): + request = FunctionRequest(function_name="echo", args=[(1, 2)]) + assert request.to_input()["serialization_format"] == "json" + assert FunctionRequest.from_input(request.to_input()).args == [[1, 2]] + + +@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": "echo", "serialization_format": fmt} + ) + + +@pytest.mark.parametrize("value", [None, False, 0, "", [], {}, (1, 2), {1: "key"}]) +def test_task_and_live_responses_share_json_conversion(value): + response = {"success": True, "json_result": value} + expected = normalize_json(value) + assert unwrap_task_response(response) == expected + assert LiveTarget("ep")._unwrap_response(response) == expected + + +@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": 42}, + ], +) +def test_task_and_live_reject_invalid_response_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.mark.parametrize("value", [b"bytes", datetime(2026, 1, 1), float("nan")]) +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)