diff --git a/CHANGELOG.md b/CHANGELOG.md index be9484adf..abc16b2f8 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/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/api/graphql.py b/runpod/api/graphql.py index 367874250..5ab5ddf27 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, @@ -110,7 +111,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 c518e8b01..19a462a0c 100644 --- a/runpod/apps/api.py +++ b/runpod/apps/api.py @@ -19,6 +19,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 = { @@ -574,6 +614,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/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..0ab02e680 100644 --- a/runpod/apps/handles.py +++ b/runpod/apps/handles.py @@ -20,6 +20,7 @@ def load_model(): ... # worker startup hook import asyncio import inspect +import json from typing import ( TYPE_CHECKING, Any, @@ -38,6 +39,7 @@ def load_model(): ... # worker startup hook from .job import Job from .markers import INIT_ATTR, is_init, route_of from .schedule import SCHEDULE_ATTR +from .serialization import require_json_value from .spec import ResourceKind, ResourceSpec, RouteSpec if TYPE_CHECKING: @@ -122,6 +124,13 @@ 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: + require_json_value( + value, "remote arguments and results must be JSON-serializable" + ) + return json.loads(json.dumps(value, allow_nan=False)) + @staticmethod async def _await_local_result(result: Any) -> Any: if inspect.isawaitable(result): @@ -139,10 +148,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 +161,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 +183,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 +245,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/serialization.py b/runpod/apps/serialization.py index 425189e1a..a7c8f5bb1 100644 --- a/runpod/apps/serialization.py +++ b/runpod/apps/serialization.py @@ -1,8 +1,8 @@ """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, provides cloudpickle codecs for explicitly +encoded values, and extracts function source for live workers. """ import ast @@ -41,6 +41,15 @@ def deserialize_result(result_b64: str) -> Any: ) from exc +def require_json_value(value: Any, message: str) -> None: + try: + import json + + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError(message) from None + + def require_json_args(fn: Callable, args: tuple, kwargs: dict) -> None: """reject arguments that cannot cross the queue data plane. @@ -48,16 +57,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..780139dc9 100644 --- a/runpod/apps/targets.py +++ b/runpod/apps/targets.py @@ -16,6 +16,7 @@ from ..user_agent import USER_AGENT from .errors import EndpointNotFound, RemoteExecutionError from .protocol import FORMAT_JSON, FunctionRequest, FunctionResponse +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,6 +694,7 @@ 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: @@ -778,6 +799,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 +819,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 +899,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 +910,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..78af5dc08 100644 --- a/runpod/apps/tasks.py +++ b/runpod/apps/tasks.py @@ -36,6 +36,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 +195,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 +208,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" @@ -312,13 +265,7 @@ async def set_timeout(self, timeout: float) -> None: async def _post(self, path: str, request: Dict[str, Any]) -> None: 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) url = f"{_proxy_url(self.pod_id)}{path}" attempts = 6 async with aiohttp.ClientSession() as session: @@ -385,9 +332,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,6 +365,111 @@ 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.""" @@ -431,13 +481,17 @@ def unwrap_task_response(response: Dict[str, Any]) -> Any: from .serialization import deserialize_result return deserialize_result(response["result"]) - return response.get("json_result") + from .serialization import require_json_value + + result = response.get("json_result") + require_json_value(result, "task results must be JSON-serializable") + return result 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 be5c18582..3782340b0 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -283,14 +283,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/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..5d0c27d57 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") 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..024d67198 100644 --- a/tests/test_apps/test_tasks.py +++ b/tests/test_apps/test_tasks.py @@ -138,7 +138,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): 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)