diff --git a/README.md b/README.md index d1e58da..a6f6199 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,33 @@ python -m runpod_sdk_runtime.bootstrap python -m runpod_sdk_runtime.task.runner ``` -`RUNPOD_RUNTIME_KIND` selects `queue` or `api` when running the shared bootstrap. +`RUNPOD_RUNTIME_KIND` selects `queue`, `api`, or deployed `task` when running the +shared bootstrap. Source-mode tasks start through `runpod_sdk_runtime.task.runner`. + +## Deployed tasks + +The control plane starts production task pods through the shared bootstrap with +`RUNPOD_RUNTIME_KIND=task`. +`RUNPOD_BUILD_ID`, `FLASH_APP_NAME`, `FLASH_ENVIRONMENT`, and `FLASH_RESOURCE_NAME` +identify the selected build and resource. The host delivers the environment's +active artifact at `/root/.runpod/artifact.tar.gz` on its read-only mount. The +bootstrap waits for that archive, extracts it and attaches its vendored +environment. The runner validates the manifest app, resource name, and +task kind, then imports the registered function and runs its init hook. + +`GET /ping` reports readiness only after initialization succeeds. Pending +initialization returns HTTP 503 with `{"ready": false}`. Bootstrap and +initialization failures return HTTP 500 with +`{"ready": false, "error_type": "BootstrapError"}`. Error details are available +through authenticated `GET /result`. Clients should treat the bootstrap failure +as terminal and retrieve its details using the task token. Task pods do not +consume serverless queue jobs. `RUNPOD_TASK_TOKEN` authenticates `/submit`, `/result`, +`/execute`, and `/timeout` on `RUNPOD_TASK_PORT` (8080 by default). + +Production execution accepts JSON `args` and `kwargs`, with an optional transport +`timeout`. Results use `json_result`, with captured `stdout` and the shared error +envelope. Requests cannot supply source, dependencies, mounts, or pickle +serialization. Source-mode tasks use the shared live executor. ## Storage bindings @@ -49,13 +75,29 @@ configuration fails startup. ## Task cleanup -Active background jobs and inline executions can run indefinitely, including -after the client disconnects. The watchdog checks for idle cleanup every 15 +Active background jobs and inline executions can run indefinitely when no +execution timeout is configured, including after the client disconnects. The +watchdog checks for idle cleanup every 15 seconds. `RUNPOD_TASK_IDLE_TIMEOUT` (600 seconds by default) applies only to pods waiting for submission or holding uncollected results. -Self-deletion requires `RUNPOD_API_KEY` and `RUNPOD_POD_ID`. The runtime exits -only after GraphQL acknowledges `podTerminate` without errors. Transient +`RUNPOD_TASK_INIT_TIMEOUT` bounds deployed task initialization, including storage +setup, module imports, and init hooks. It defaults to 600 seconds and must be +finite and positive. Initialization timeout produces a terminal `BootstrapError` +and triggers pod deletion, with retries until deletion is acknowledged. Client +contact cannot extend this deadline, and completion after timeout cannot make +the pod ready. Idle monitoring begins after initialization succeeds or fails; +the initialization deadline is independent of the execution timeout. + +`RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS` is an optional, finite positive production +policy cap selected by the control plane. Its deadline starts at the first accepted +execution. Request timeouts can shorten that deadline but cannot extend it. The +watchdog enforces execution deadlines even while work is active. + +Self-deletion requires `RUNPOD_API_KEY` and `RUNPOD_POD_ID`. Production task pods +call `flashTaskTerminate` with the internal pod key in `Authorization` and the pod ID +in `X-RunPod-Pod-Id`; source-mode pods call `podTerminate`. The runtime exits +only after GraphQL acknowledges deletion without errors. Transient failures receive three bounded attempts; any failed deletion leaves the runtime alive for the next watchdog pass. Container exit alone does not stop pod billing. diff --git a/src/runpod_sdk_runtime/bootstrap.py b/src/runpod_sdk_runtime/bootstrap.py index f0f5857..859f583 100644 --- a/src/runpod_sdk_runtime/bootstrap.py +++ b/src/runpod_sdk_runtime/bootstrap.py @@ -27,6 +27,8 @@ job-take loop answers every job with a structured BootstrapError so the failure surfaces in job responses (queue endpoints), or an http server answers every request with the error (api endpoints). +task pods serve readiness and execution errors over their token-protected +http transport; their init hooks run before readiness succeeds. stdlib only: this file runs on the image's python before any artifact environment exists. runtime images bake it in, and custom images install @@ -57,11 +59,13 @@ WORKER_MODULES = { "queue": "runpod_sdk_runtime.queue.worker", "api": "runpod_sdk_runtime.api.server", + "task": "runpod_sdk_runtime.task.runner", } SDK_MODULES = { "queue": ("runpod.apps", "runpod.serverless"), "api": ("runpod.apps", "fastapi", "uvicorn"), + "task": ("runpod.apps", "runpod.apps.handles"), } @@ -87,6 +91,41 @@ def _resource_name(): ) +def _validate_task_environment(): + if _runtime_kind() != "task": + raise PhaseError("identity", "RUNPOD_RUNTIME_KIND must be task") + for key in ( + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_BUILD_ID", + ): + if not os.environ.get(key, "").strip(): + raise PhaseError("identity", f"{key} is required for deployed tasks") + + +def _task_resource(manifest): + if not isinstance(manifest, dict): + raise PhaseError("attach", "manifest must be an object") + if manifest.get("app") != os.environ.get("FLASH_APP_NAME"): + raise PhaseError("attach", "manifest app does not match FLASH_APP_NAME") + resources = manifest.get("resources") + if not isinstance(resources, list) or any( + not isinstance(entry, dict) for entry in resources + ): + raise PhaseError("attach", "manifest resources must be an array of objects") + entries = [r for r in resources if r.get("name") == _resource_name()] + if len(entries) != 1 or entries[0].get("kind") != "task": + raise PhaseError("attach", "manifest must contain exactly one matching task") + entry = entries[0] + if any( + not isinstance(entry.get(key), str) or not entry[key].strip() + for key in ("module", "qualname") + ): + raise PhaseError("attach", "task module and qualname are required") + return entry + + # ---------------------------------------------------------------- locate @@ -175,8 +214,11 @@ def _manifest(app_dir): path = os.path.join(app_dir, MANIFEST_NAME) if not os.path.isfile(path): raise PhaseError("attach", f"manifest not found at {path}") - with open(path) as f: - return json.load(f) + try: + with open(path) as f: + return json.load(f) + except (OSError, ValueError) as exc: + raise PhaseError("attach", f"failed to read manifest: {exc}") from exc def _sdk_importable(paths): @@ -396,7 +438,11 @@ def log_message(self, *args): def _report_error(error): _log(f"FAILED {error}") - if _runtime_kind() == "api": + if _runtime_kind() == "task": + from runpod_sdk_runtime.task.runner import main as task_main + + task_main(bootstrap_error=_error_payload(error)) + elif _runtime_kind() == "api": _api_error_server(error) else: _queue_error_loop(error) @@ -407,10 +453,15 @@ def _report_error(error): def main(): try: + task = _runtime_kind() == "task" + if task: + _validate_task_environment() if _resource_name(): app_dir = _locate() paths = _attach(app_dir) manifest = _manifest(app_dir) + if task: + _task_resource(manifest) _verify_excluded(manifest) _install_system(manifest) ok, err = _sdk_importable(paths) @@ -426,6 +477,10 @@ def main(): _serve([]) except PhaseError as exc: _report_error(exc) + except Exception as exc: + if not task: + raise + _report_error(PhaseError("bootstrap", str(exc))) if __name__ == "__main__": diff --git a/src/runpod_sdk_runtime/executor.py b/src/runpod_sdk_runtime/executor.py index 516a0d8..ade184c 100644 --- a/src/runpod_sdk_runtime/executor.py +++ b/src/runpod_sdk_runtime/executor.py @@ -102,12 +102,12 @@ def _deserialize_args(request): return args, kwargs -def _serialize_result(result, fmt): +def _serialize_result(result, fmt, *, allow_nan=True): if fmt == "json": # json-format requests promise json on the wire; failing here # keeps dev behavior identical to a deployed endpoint try: - json.dumps(result) + json.dumps(result, allow_nan=allow_nan) except (TypeError, ValueError) as exc: raise TypeError( f"return value must be json-serializable " @@ -177,13 +177,16 @@ def serialize_chunk(chunk, request): return response -def execute_request(request): +def execute_request(request, *, prepared=None, run_awaitable=None): """run one FunctionRequest to completion, returning a response dict.""" stdout_io = io.StringIO() + strict_json = prepared is not None + await_result = run_awaitable or _run_awaitable try: - prepared, error_response = resolve_request(request) - if error_response is not None: - return error_response + if prepared is None: + prepared, error_response = resolve_request(request) + if error_response is not None: + return error_response fn, args, kwargs = prepared # stdout is teed: captured for the job response and written @@ -194,18 +197,20 @@ def execute_request(request): with redirect_stdout(_Tee(stdout_io, sys.__stdout__)): result = fn(*args, **kwargs) if hasattr(result, "__await__"): - result = _run_awaitable(result) + result = await_result(result) # generators aggregate so live .remote() matches the # deployed worker's return_aggregate_stream output elif inspect.isgenerator(result): result = list(result) elif inspect.isasyncgen(result): - result = _run_awaitable(_drain_async_gen(result)) + result = await_result(_drain_async_gen(result)) response = {"success": True, "stdout": stdout_io.getvalue()} response.update( _serialize_result( - result, request.get("serialization_format", "cloudpickle") + result, + request.get("serialization_format", "cloudpickle"), + allow_nan=not strict_json, ) ) return response diff --git a/src/runpod_sdk_runtime/task/deployed.py b/src/runpod_sdk_runtime/task/deployed.py new file mode 100644 index 0000000..97d5108 --- /dev/null +++ b/src/runpod_sdk_runtime/task/deployed.py @@ -0,0 +1,117 @@ +"""registered task loading and json-only execution.""" + +import asyncio +import importlib +import inspect +import json +import os +import sys +import threading + +from runpod_sdk_runtime import bootstrap +from runpod_sdk_runtime.executor import _await, execute_request + +_loop = None +_loop_thread = None +_loop_lock = threading.Lock() + + +def _serve_loop(loop): + asyncio.set_event_loop(loop) + try: + loop.run_forever() + finally: + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + loop.close() + + +def _run_awaitable(awaitable): + global _loop, _loop_thread + with _loop_lock: + if _loop is None: + _loop = asyncio.new_event_loop() + _loop_thread = threading.Thread( + target=_serve_loop, args=(_loop,), daemon=True + ) + _loop_thread.start() + future = asyncio.run_coroutine_threadsafe(_await(awaitable), _loop) + return future.result() + + +def close(): + global _loop, _loop_thread + with _loop_lock: + if _loop is None: + return + _loop.call_soon_threadsafe(_loop.stop) + _loop_thread.join() + _loop = None + _loop_thread = None + + +async def _initialize(handle): + result = handle._init_fn() + if inspect.isawaitable(result): + await result + + +def load_handle(): + bootstrap._validate_task_environment() + app_dir = os.environ.get("RUNPOD_APP_DIR", "/app") + entry = bootstrap._task_resource(bootstrap._manifest(app_dir)) + if app_dir not in sys.path: + sys.path.insert(0, app_dir) + module = importlib.import_module(entry["module"]) + + from runpod.apps.handles import FunctionHandle + from runpod.apps.spec import ResourceKind + + target = module + for name in entry["qualname"].split("."): + target = getattr(target, name, None) + handles = { + value + for value in [target, *vars(module).values()] + if isinstance(value, FunctionHandle) + and value.spec.name == entry["name"] + and value.spec.kind == ResourceKind.TASK + and value._app.name == os.environ["FLASH_APP_NAME"] + and value._fn.__qualname__ == entry["qualname"] + } + if len(handles) != 1: + raise RuntimeError("module must expose exactly one matching @app.task handle") + handle = handles.pop() + if handle._init_fn is not None: + _run_awaitable(_initialize(handle)) + return handle + + +def validate_request(request, *, timeout_only=False): + allowed = ( + {"timeout"} + if timeout_only + else {"args", "kwargs", "serialization_format", "timeout"} + ) + if request.keys() - allowed: + raise ValueError("deployed tasks accept only json args, kwargs and timeout") + if request.get("serialization_format", "json") != "json": + raise ValueError("deployed tasks require json serialization") + if not isinstance(request.get("args", []), list): + raise ValueError("args must be a json array") + if not isinstance(request.get("kwargs", {}), dict): + raise ValueError("kwargs must be a json object") + json.dumps(request, allow_nan=False) + + +def execute(handle, request): + return execute_request( + {"serialization_format": "json"}, + prepared=(handle._fn, request.get("args", []), request.get("kwargs", {})), + run_awaitable=_run_awaitable, + ) diff --git a/src/runpod_sdk_runtime/task/runner.py b/src/runpod_sdk_runtime/task/runner.py index efe9cc8..72d98e1 100644 --- a/src/runpod_sdk_runtime/task/runner.py +++ b/src/runpod_sdk_runtime/task/runner.py @@ -31,6 +31,7 @@ from runpod_sdk_runtime.executor import execute_request from runpod_sdk_runtime.mounts import configure_mounts +from runpod_sdk_runtime.task import deployed PORT = int(os.environ.get("RUNPOD_TASK_PORT", "8080")) TOKEN = os.environ.get("RUNPOD_TASK_TOKEN", "") @@ -49,6 +50,11 @@ _inline_executions = {"count": 0} _execution_deadline = None _watchdog_wakeup = threading.Event() +_ready = True +_bootstrap_error = None +_deployed_handle = None +_timeout_cap = None +_production = False def _touch_contact(): @@ -66,6 +72,30 @@ def _set_timeout(timeout): _watchdog_wakeup.set() +def _start_execution(timeout): + _set_timeout(_timeout_cap) + _set_timeout(timeout) + + +def _configured_timeout(key="RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", default=None): + value = os.environ.get(key) + if value is None: + return default + try: + timeout = float(value) + except ValueError: + timeout = float("nan") + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError(f"{key} must be finite and positive") + return timeout + + +def _execute(request): + if _production: + return deployed.execute(_deployed_handle, request) + return execute_request(request) + + def _should_self_terminate(status, last_contact, now, idle_timeout, inline=0): if status == "RUNNING" or inline > 0: return False @@ -84,11 +114,14 @@ def _terminate_self(): ) return False api_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") + task = os.environ.get("RUNPOD_RUNTIME_KIND") == "task" + mutation = "flashTaskTerminate" if task else "podTerminate" + input_type = "FlashTaskTerminateInput" if task else "PodTerminateInput" payload = json.dumps( { "query": ( - "mutation podTerminate($input: PodTerminateInput!) " - "{ podTerminate(input: $input) }" + f"mutation {mutation}($input: {input_type}!) " + f"{{ {mutation}(input: $input) }}" ), "variables": {"input": {"podId": pod_id}}, } @@ -101,6 +134,8 @@ def _terminate_self(): "Authorization": f"Bearer {api_key}", }, ) + if task: + request.add_header("X-RunPod-Pod-Id", pod_id) for attempt in range(TERMINATE_ATTEMPTS): retryable = False try: @@ -111,8 +146,8 @@ def _terminate_self(): if ( not errors and isinstance(data, dict) - and "podTerminate" in data - and data["podTerminate"] is None + and mutation in data + and (data[mutation] is True if task else data[mutation] is None) ): sys.stderr.write(f"[task-runner] pod {pod_id} deletion acknowledged\n") os._exit(0) @@ -191,16 +226,31 @@ def _authed(self): def _read_request(self): length = int(self.headers.get("Content-Length", "0")) + if length < 0: + raise ValueError("invalid content length") return json.loads(self.rfile.read(length) or b"{}") + def _available(self, *, public=False): + if _bootstrap_error is not None: + error = {"error_type": "BootstrapError"} if public else _bootstrap_error + self._send(500, {"ready": False, **error}) + return False + if not _ready: + self._send(503, {"ready": False}) + return False + return True + def do_GET(self): # noqa: N802 - BaseHTTPRequestHandler api if self.path == "/ping": - self._send(200, {"ready": True}) + if self._available(public=True): + self._send(200, {"ready": True}) return if self.path == "/result": if not self._authed(): self._send(401, {"error": "unauthorized"}) return + if not self._available(): + return with _job_lock: self._send( 200, @@ -219,21 +269,32 @@ def do_POST(self): # noqa: N802 - BaseHTTPRequestHandler api if self.path not in ("/execute", "/submit", "/timeout"): self._send(404, {"error": "not found"}) return + if not self._available(): + return try: request = self._read_request() - except (ValueError, TypeError): + except (ValueError, TypeError, OverflowError): self._send(400, {"error": "invalid request"}) return if not isinstance(request, dict): self._send(400, {"error": "invalid request"}) return + if _production: + try: + deployed.validate_request(request, timeout_only=self.path == "/timeout") + except (ValueError, TypeError) as exc: + self._send(400, {"error": str(exc)}) + return 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 - ): + try: + 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 + except (ValueError, OverflowError): self._send(400, {"error": "timeout must be a finite non-negative number"}) return if self.path == "/timeout": @@ -243,25 +304,35 @@ def do_POST(self): # noqa: N802 - BaseHTTPRequestHandler api return if self.path == "/execute": with _job_lock: - _set_timeout(timeout) + if _production and ( + _job_state["status"] != "NONE" or _inline_executions["count"] + ): + self._send(409, {"error": "this task has already been submitted"}) + return + _start_execution(timeout) _inline_executions["count"] += 1 try: - self._send(200, execute_request(request)) + self._send(200, _execute(request)) finally: with _job_lock: _inline_executions["count"] -= 1 + if _production: + _job_state["status"] = "DONE" return if self.path == "/submit": with _job_lock: - if _job_state["status"] == "RUNNING": + if _job_state["status"] == "RUNNING" or ( + _production + and (_job_state["status"] != "NONE" or _inline_executions["count"]) + ): self._send(409, {"error": "a job is already running"}) return - _set_timeout(timeout) + _start_execution(timeout) _job_state["status"] = "RUNNING" _job_state["response"] = None def run(): - response = execute_request(request) + response = _execute(request) with _job_lock: _job_state["status"] = "DONE" _job_state["response"] = response @@ -279,15 +350,102 @@ def log_message(self, format, *args): # noqa: A002 - stdlib signature pass -def main() -> None: +def main(bootstrap_error=None) -> None: + global _ready, _bootstrap_error, _deployed_handle, _timeout_cap + global _production if not TOKEN: sys.stderr.write("[task-runner] RUNPOD_TASK_TOKEN not set, exiting\n") sys.exit(1) - configure_mounts() - _touch_contact() - threading.Thread(target=_watchdog, daemon=True).start() + _bootstrap_error = bootstrap_error + _ready = False + _deployed_handle = None + _timeout_cap = None + _production = os.environ.get("RUNPOD_RUNTIME_KIND") == "task" or any( + key in os.environ + for key in ( + "RUNPOD_BUILD_ID", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_RESOURCE_NAME", + ) + ) + + def start_watchdog(): + _touch_contact() + threading.Thread(target=_watchdog, daemon=True).start() + + initialization_done = threading.Event() + + def watch_initialization(): + global _bootstrap_error + while True: + remaining = init_deadline - time.monotonic() + if remaining <= 0: + break + if initialization_done.wait(min(remaining, WATCHDOG_INTERVAL)): + return + with _job_lock: + if initialization_done.is_set(): + return + _bootstrap_error = { + "error_type": "BootstrapError", + "error_message": "task initialization timed out", + } + initialization_done.set() + while not _terminate_self(): + time.sleep(WATCHDOG_INTERVAL) + + def initialize(): + global _ready, _bootstrap_error, _deployed_handle, _timeout_cap + handle = None + error = None + try: + _timeout_cap = _configured_timeout() + configure_mounts() + handle = deployed.load_handle() + except Exception as exc: + error = { + "error_type": "BootstrapError", + "error_message": f"task initialization failed: {exc}", + } + with _job_lock: + if initialization_done.is_set() or time.monotonic() >= init_deadline: + return + _deployed_handle = handle + _bootstrap_error = error + _ready = error is None + initialization_done.set() + start_watchdog() + + if not _production and bootstrap_error is None: + configure_mounts() + _ready = True server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler) - server.serve_forever() + initializer = None + try: + if _production and bootstrap_error is None: + try: + init_timeout = _configured_timeout("RUNPOD_TASK_INIT_TIMEOUT", 600.0) + except ValueError as exc: + _bootstrap_error = { + "error_type": "BootstrapError", + "error_message": f"task initialization failed: {exc}", + } + start_watchdog() + else: + init_deadline = time.monotonic() + init_timeout + threading.Thread(target=watch_initialization, daemon=True).start() + initializer = threading.Thread(target=initialize, daemon=True) + initializer.start() + else: + start_watchdog() + server.serve_forever() + finally: + if initializer is not None: + initializer.join() + server.server_close() + deployed.close() if __name__ == "__main__": diff --git a/tests/test_deployed_task.py b/tests/test_deployed_task.py new file mode 100644 index 0000000..762f44a --- /dev/null +++ b/tests/test_deployed_task.py @@ -0,0 +1,912 @@ +"""production task bootstrap, http transport and lifecycle contracts.""" + +import io +import json +import threading +import time +import urllib.error +import urllib.request +from http.server import ThreadingHTTPServer +from unittest.mock import Mock + +import pytest + +from runpod_sdk_runtime import bootstrap +from runpod_sdk_runtime.task import deployed, runner + +TOKEN = "execution-token" + + +@pytest.fixture +def artifact(tmp_path, monkeypatch): + import hashlib + + module = "task_" + hashlib.sha256(str(tmp_path).encode()).hexdigest()[:12] + source = """import runpod +app = runpod.App("example") +initialized = False +@app.task(name="work", cpu="cpu3c") +async def work(value=1, *, mode="success"): + assert initialized + print("task output") + if mode == "error": + raise ValueError("task failed") + if mode == "nonfinite": + return float("nan") + if mode == "nonjson": + return {1, 2} + return {"value": value} +@work.init +async def initialize(): + global initialized + initialized = True +""" + (tmp_path / (module + ".py")).write_text(source) + manifest = { + "app": "example", + "resources": [ + {"name": "work", "kind": "task", "module": module, "qualname": "work"} + ], + } + (tmp_path / bootstrap.MANIFEST_NAME).write_text(json.dumps(manifest)) + for key, value in { + "RUNPOD_RUNTIME_KIND": "task", + "RUNPOD_BUILD_ID": "build-1", + "FLASH_APP_NAME": "example", + "FLASH_ENVIRONMENT": "production", + "FLASH_RESOURCE_NAME": "work", + "FLASH_BUILD_ARTIFACT_PATH": str(tmp_path / "artifact.tar.gz"), + "RUNPOD_APP_DIR": str(tmp_path), + "RUNPOD_MOUNTS": "[]", + }.items(): + monkeypatch.setenv(key, value) + monkeypatch.delenv("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", raising=False) + monkeypatch.syspath_prepend(str(tmp_path)) + return tmp_path, manifest + + +def request(url, method="GET", body=None, token=TOKEN): + headers = {"Authorization": f"Bearer {token}"} if token else {} + data = json.dumps(body).encode() if body is not None else None + req = urllib.request.Request(url, method=method, headers=headers, data=data) + try: + response = urllib.request.urlopen(req, timeout=5) + except urllib.error.HTTPError as exc: + response = exc + with response: + return response.status, json.load(response) + + +def wait_ready(url): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + status, body = request(url + "/ping", token=None) + if status != 503: + return status, body + time.sleep(0.01) + pytest.fail("task did not initialize") + + +@pytest.fixture +def start_server(artifact, monkeypatch): + from runpod.apps.volume import _configure_mounts + + servers = [] + monkeypatch.setattr(runner, "TOKEN", "execution-token") + monkeypatch.setattr(runner, "_watchdog", lambda: None) + for attr, value in { + "_job_state": {"status": "NONE", "response": None}, + "_execution_deadline": None, + "_watchdog_wakeup": threading.Event(), + "_inline_executions": {"count": 0}, + "_ready": False, + "_bootstrap_error": None, + "_deployed_handle": None, + "_timeout_cap": None, + "_production": False, + }.items(): + monkeypatch.setattr(runner, attr, value) + + def start(error=None): + httpd = ThreadingHTTPServer(("127.0.0.1", 0), runner.Handler) + httpd.daemon_threads = True + monkeypatch.setattr(runner, "ThreadingHTTPServer", lambda *_: httpd) + thread = threading.Thread( + target=runner.main, kwargs={"bootstrap_error": error}, daemon=True + ) + servers.append((httpd, thread)) + thread.start() + return f"http://127.0.0.1:{httpd.server_port}" + + yield start + for httpd, thread in servers: + httpd.shutdown() + thread.join(5) + httpd.server_close() + _configure_mounts([]) + + +def test_registered_task_submit_result_and_single_shot(start_server): + url = start_server() + assert wait_ready(url) == (200, {"ready": True}) + assert request(url + "/submit", "POST", {"args": [42]}) == ( + 200, + {"status": "RUNNING"}, + ) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + _, result = request(url + "/result") + if result["status"] == "DONE": + break + time.sleep(0.01) + assert result == { + "status": "DONE", + "response": { + "success": True, + "json_result": {"value": 42}, + "stdout": "task output\n", + }, + } + assert request(url + "/submit", "POST", {})[0] == 409 + + +@pytest.mark.parametrize( + "mode, message", + [ + ("error", "task failed"), + ("nonfinite", "Out of range float"), + ("nonjson", "json-serializable"), + ], +) +def test_registered_task_errors_preserve_stdout(start_server, mode, message): + url = start_server() + assert wait_ready(url)[0] == 200 + status, response = request(url + "/execute", "POST", {"kwargs": {"mode": mode}}) + assert status == 200 + assert response["success"] is False + assert message in response["error"] + assert response["stdout"] == "task output\n" + + +@pytest.mark.parametrize( + "body", + [ + {"function_code": "raise RuntimeError()"}, + {"function_name": "other"}, + {"dependencies": []}, + {"system_dependencies": []}, + {"mounts": []}, + {"serialization_format": "cloudpickle"}, + {"args": "pickle"}, + {"kwargs": []}, + {"args": [float("nan")]}, + {"args": None}, + {"kwargs": None}, + {"timeout": 10**400}, + ], +) +def test_production_rejects_request_overrides(start_server, body): + url = start_server() + assert wait_ready(url)[0] == 200 + assert request(url + "/submit", "POST", body)[0] == 400 + assert runner._job_state["status"] == "NONE" + assert runner._execution_deadline is None + + +@pytest.mark.parametrize("body", [[], None, "text", 4]) +def test_production_rejects_malformed_body(start_server, body): + url = start_server() + assert wait_ready(url)[0] == 200 + req = urllib.request.Request( + url + "/submit", + data=json.dumps(body).encode(), + headers={"Authorization": "Bearer execution-token"}, + ) + with pytest.raises(urllib.error.HTTPError) as exc: + urllib.request.urlopen(req, timeout=5) + assert exc.value.code == 400 + exc.value.close() + + +def test_production_readiness_waits_for_init(start_server, monkeypatch): + entered, release = threading.Event(), threading.Event() + original = deployed.load_handle + + def load(): + entered.set() + release.wait(5) + return original() + + monkeypatch.setattr(deployed, "load_handle", load) + url = start_server() + try: + assert entered.wait(5) + assert request(url + "/ping", token=None) == (503, {"ready": False}) + assert request(url + "/submit", "POST", {})[0] == 503 + finally: + release.set() + assert wait_ready(url)[0] == 200 + + +def test_init_failure_has_task_error_surface(start_server, monkeypatch): + monkeypatch.setattr(deployed, "load_handle", Mock(side_effect=ValueError("init"))) + url = start_server() + status, body = wait_ready(url) + assert status == 500 + assert body["ready"] is False + assert body["error_type"] == "BootstrapError" + assert "error_message" not in body + status, body = request(url + "/result") + assert status == 500 + assert "init" in body["error_message"] + assert request(url + "/submit", "POST", {}) == (status, body) + assert request(url + "/result") == (status, body) + assert request(url + "/submit", "POST", {}, token="wrong")[0] == 401 + + +@pytest.mark.parametrize( + "key,value", + [ + ("RUNPOD_RUNTIME_KIND", "queue"), + ("FLASH_APP_NAME", ""), + ("FLASH_RESOURCE_NAME", ""), + ("FLASH_ENVIRONMENT", ""), + ("RUNPOD_BUILD_ID", ""), + ], +) +def test_invalid_production_identity_never_falls_back( + artifact, monkeypatch, key, value +): + monkeypatch.setenv(key, value) + with pytest.raises(bootstrap.PhaseError): + deployed.load_handle() + + +@pytest.mark.parametrize("change", ["app", "kind", "name", "qualname", "module"]) +def test_manifest_must_identify_registered_task(artifact, change): + path, manifest = artifact + if change == "app": + manifest["app"] = "other" + else: + manifest["resources"][0][change] = "other" + (path / bootstrap.MANIFEST_NAME).write_text(json.dumps(manifest)) + with pytest.raises((bootstrap.PhaseError, RuntimeError, ImportError)): + deployed.load_handle() + + +def test_duplicate_task_manifest_rejected(artifact): + _, manifest = artifact + manifest["resources"] *= 2 + with pytest.raises(bootstrap.PhaseError, match="exactly one"): + bootstrap._task_resource(manifest) + + +def test_manifest_qualname_resolves_class_attribute(artifact): + path, manifest = artifact + entry = manifest["resources"][0] + entry["qualname"] = "Tasks.work" + (path / bootstrap.MANIFEST_NAME).write_text(json.dumps(manifest)) + (path / (entry["module"] + ".py")).write_text( + 'import runpod\napp = runpod.App("example")\nclass Tasks:\n' + ' @app.task(name="work", cpu="cpu3c")\n' + " def work():\n return 7\n" + ) + assert deployed.execute(deployed.load_handle(), {})["json_result"] == 7 + + +def test_production_mode_cannot_fall_back_after_startup(start_server, monkeypatch): + url = start_server() + assert wait_ready(url)[0] == 200 + monkeypatch.delenv("RUNPOD_RUNTIME_KIND") + assert ( + request( + url + "/execute", + "POST", + { + "function_name": "live", + "function_code": "def live(): return 7", + "serialization_format": "json", + }, + )[0] + == 400 + ) + assert request(url + "/execute", "POST", {})[1]["json_result"] == {"value": 1} + + +def test_configured_cap_starts_at_first_execution(start_server, monkeypatch): + monkeypatch.setenv("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", "60") + url = start_server() + assert wait_ready(url)[0] == 200 + assert runner._execution_deadline is None + start = time.monotonic() + assert request(url + "/execute", "POST", {"timeout": 600})[0] == 200 + deadline = runner._execution_deadline + assert start + 60 <= deadline <= time.monotonic() + 60 + assert request(url + "/timeout", "POST", {"timeout": 6000})[0] == 200 + assert request(url + "/timeout", "POST", {})[0] == 200 + assert runner._execution_deadline == deadline + request(url + "/timeout", "POST", {"timeout": 1}) + assert runner._execution_deadline < deadline + + +@pytest.mark.parametrize("value", ["", "0", "-1", "nan", "inf", "text"]) +def test_invalid_configured_cap_fails_readiness(start_server, monkeypatch, value): + monkeypatch.setenv("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", value) + url = start_server() + status, body = wait_ready(url) + assert status == 500 + assert body == {"ready": False, "error_type": "BootstrapError"} + status, body = request(url + "/result") + assert status == 500 + assert "finite and positive" in body["error_message"] + + +def test_no_policy_cap_does_not_add_deadline(start_server): + url = start_server() + assert wait_ready(url)[0] == 200 + request(url + "/execute", "POST", {}) + assert runner._execution_deadline is None + + +def test_deployed_execution_never_resolves_source_or_pickle(artifact, monkeypatch): + from runpod_sdk_runtime import executor + + resolve = Mock(side_effect=AssertionError("source resolution")) + pickle = Mock(side_effect=AssertionError("pickle resolution")) + monkeypatch.setattr(executor, "resolve_request", resolve) + monkeypatch.setattr(executor, "_load_cloudpickle", pickle) + result = deployed.execute(deployed.load_handle(), {"kwargs": {"value": 7}}) + assert result["json_result"] == {"value": 7} + resolve.assert_not_called() + pickle.assert_not_called() + + +def test_source_execution_ignores_production_policy_cap(start_server, monkeypatch): + for key in ( + "RUNPOD_RUNTIME_KIND", + "RUNPOD_BUILD_ID", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + ): + monkeypatch.delenv(key, raising=False) + monkeypatch.setenv("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", "invalid") + url = start_server() + assert wait_ready(url)[0] == 200 + status, response = request( + url + "/execute", + "POST", + { + "function_name": "live", + "function_code": "def live(): return 7", + "serialization_format": "json", + }, + ) + assert status == 200 + assert response["json_result"] == 7 + assert runner._execution_deadline is None + + +@pytest.mark.parametrize( + "result,expected", [(True, True), (False, False), (None, False)] +) +def test_production_self_termination_uses_pod_auth(monkeypatch, result, expected): + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "task") + monkeypatch.setenv("RUNPOD_POD_ID", "pod-1") + monkeypatch.setenv("RUNPOD_API_KEY", "internal-key") + response = io.BytesIO(json.dumps({"data": {"flashTaskTerminate": result}}).encode()) + urlopen = Mock(return_value=response) + exits = [] + monkeypatch.setattr(runner.urllib.request, "urlopen", urlopen) + monkeypatch.setattr(runner.os, "_exit", exits.append) + assert runner._terminate_self() is expected + req = urlopen.call_args.args[0] + assert req.get_header("Authorization") == "Bearer internal-key" + assert req.get_header("X-runpod-pod-id") == "pod-1" + payload = json.loads(req.data) + assert "FlashTaskTerminateInput!" in payload["query"] + assert "flashTaskTerminate(input:" in payload["query"] + assert payload["variables"] == {"input": {"podId": "pod-1"}} + assert exits == ([0] if expected else []) + + +def test_task_bootstrap_error_never_enters_queue_loop(monkeypatch): + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "task") + main = Mock() + queue = Mock() + monkeypatch.setattr(runner, "main", main) + monkeypatch.setattr(bootstrap, "_queue_error_loop", queue) + bootstrap._report_error(bootstrap.PhaseError("locate", "missing")) + queue.assert_not_called() + assert main.call_args.kwargs["bootstrap_error"]["error_type"] == "BootstrapError" + + +def test_production_bootstrap_uses_prebuilt_app(artifact, monkeypatch): + path, _ = artifact + prebuilt = path / "prebuilt" + prebuilt.mkdir() + monkeypatch.setattr(bootstrap, "PREBUILT_APP_DIR", str(prebuilt)) + monkeypatch.setattr(bootstrap, "ARTIFACT_PATH", str(path / "artifact.tar.gz")) + monkeypatch.setattr(bootstrap, "APP_DIR", str(path / "app")) + monkeypatch.setattr(bootstrap, "ARTIFACT_WAIT_SECONDS", 0) + assert bootstrap._locate() == str(prebuilt) + + +def test_production_bootstrap_reuses_unpacked_app(artifact, monkeypatch): + path, _ = artifact + (path / bootstrap.UNPACK_MARKER).write_text(str(time.time())) + monkeypatch.setattr(bootstrap, "APP_DIR", str(path)) + assert bootstrap._locate() == str(path) + + +def test_task_identity_accepts_host_artifact_path(artifact, monkeypatch): + monkeypatch.delenv("FLASH_BUILD_ARTIFACT_PATH") + bootstrap._validate_task_environment() + + +@pytest.mark.parametrize("contents", ["{bad", "[]", '{"app":"wrong"}']) +def test_malformed_manifest_reports_bootstrap_error(artifact, monkeypatch, contents): + path, _ = artifact + (path / bootstrap.MANIFEST_NAME).write_text(contents) + monkeypatch.setattr(bootstrap, "_locate", lambda: str(path)) + report = Mock() + monkeypatch.setattr(bootstrap, "_report_error", report) + bootstrap.main() + assert report.call_args.args[0].phase == "attach" + + +@pytest.mark.parametrize("async_generator", [True, False]) +def test_registered_generators_aggregate_json(artifact, async_generator): + path, manifest = artifact + module = manifest["resources"][0]["module"] + prefix = "async " if async_generator else "" + (path / (module + ".py")).write_text( + 'import runpod\napp = runpod.App("example")\n' + '@app.task(name="work", cpu="cpu3c")\n' + f'{prefix}def work():\n print("chunks")\n yield 1\n yield 2\n' + ) + response = deployed.execute(deployed.load_handle(), {}) + assert response == {"success": True, "json_result": [1, 2], "stdout": "chunks\n"} + + +def test_real_bootstrap_archive_http_and_deadline_cleanup( + artifact, tmp_path, monkeypatch +): + import os + import socket + import subprocess + import sys + import tarfile + from http.server import BaseHTTPRequestHandler + + path, manifest = artifact + (path / "env").mkdir() + (path / "env" / "vendored_probe.py").write_text("value = 42\n") + module = manifest["resources"][0]["module"] + source = path / (module + ".py") + source.write_text("import vendored_probe\n" + source.read_text()) + archive = path / "artifact.tar.gz" + with tarfile.open(archive, "w:gz") as tar: + for name in (source.name, bootstrap.MANIFEST_NAME, "env"): + tar.add(path / name, arcname=name) + + terminated = threading.Event() + calls = [] + + class GraphQL(BaseHTTPRequestHandler): + def do_POST(self): # noqa: N802 + calls.append( + ( + dict(self.headers), + json.load( + io.BytesIO(self.rfile.read(int(self.headers["Content-Length"]))) + ), + ) + ) + body = b'{"data":{"flashTaskTerminate":true}}' + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + terminated.set() + + def log_message(self, *_): + pass + + graphql = ThreadingHTTPServer(("127.0.0.1", 0), GraphQL) + graphql_thread = threading.Thread(target=graphql.serve_forever, daemon=True) + graphql_thread.start() + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + env = dict(os.environ) + env.update( + { + "RUNPOD_TASK_PORT": str(port), + "RUNPOD_TASK_TOKEN": "execution-token", + "RUNPOD_POD_ID": "pod-integration", + "RUNPOD_API_KEY": "internal-integration-key", + "RUNPOD_API_BASE_URL": f"http://127.0.0.1:{graphql.server_port}", + "RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS": "1.5", + "RUNPOD_APP_DIR": str(tmp_path / "unpacked"), + } + ) + with (tmp_path / "worker.log").open("w+") as log: + process = subprocess.Popen( # noqa: S603 + [sys.executable, "-m", "runpod_sdk_runtime.bootstrap"], + env=env, + stdout=log, + stderr=log, + ) + try: + url = f"http://127.0.0.1:{port}" + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + try: + if request(url + "/ping")[0] == 200: + break + except OSError: + pass + assert process.poll() is None + time.sleep(0.02) + else: + pytest.fail("bootstrap did not become ready") + assert not terminated.is_set() + status, response = request(url + "/execute", "POST", {"args": [42]}) + assert status == 200 + assert response["json_result"] == {"value": 42} + assert response["stdout"] == "task output\n" + assert terminated.wait(5) + assert process.wait(timeout=5) == 0 + headers, payload = calls[0] + assert headers["Authorization"] == "Bearer internal-integration-key" + assert headers["X-Runpod-Pod-Id"] == "pod-integration" + assert "flashTaskTerminate" in payload["query"] + assert payload["variables"]["input"] == {"podId": "pod-integration"} + log.seek(0) + output = log.read() + assert "internal-integration-key" not in output + assert "execution-token" not in output + finally: + if process.poll() is None: + process.terminate() + process.wait(timeout=5) + graphql.shutdown() + graphql_thread.join(5) + graphql.server_close() + + +@pytest.mark.parametrize("kind", [None, "", "queue", "api", "invalid"]) +@pytest.mark.parametrize( + "marker", + [ + "RUNPOD_BUILD_ID", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_RESOURCE_NAME", + ], +) +def test_runner_rejects_partial_deployment_identity( + start_server, monkeypatch, kind, marker +): + for key in ( + "RUNPOD_RUNTIME_KIND", + "RUNPOD_BUILD_ID", + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_RESOURCE_NAME", + ): + monkeypatch.delenv(key, raising=False) + monkeypatch.setenv(marker, "deployment") + if kind is not None: + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", kind) + url = start_server() + status, body = wait_ready(url) + assert status == 500 + assert body["error_type"] == "BootstrapError" + assert "error_message" not in body + status, body = request(url + "/result") + assert status == 500 + assert "RUNPOD_RUNTIME_KIND must be task" in body["error_message"] + for endpoint in ("/execute", "/submit"): + assert request( + url + endpoint, + "POST", + { + "function_name": "live", + "function_code": "def live(): return 7", + "serialization_format": "json", + }, + ) == (status, body) + + +@pytest.mark.parametrize("endpoint", ["/execute", "/submit"]) +@pytest.mark.parametrize("generator", [False, True]) +def test_init_resources_survive_until_async_execution( + artifact, start_server, endpoint, generator +): + path, manifest = artifact + module = manifest["resources"][0]["module"] + result = "yield value" if generator else "return value" + (path / (module + ".py")).write_text( + 'import asyncio\nimport runpod\napp = runpod.App("example")\n' + '@app.task(name="work", cpu="cpu3c")\n' + "async def work():\n" + " assert asyncio.get_running_loop() is loop\n" + " gate.set()\n" + " value = await asyncio.wait_for(future, 1)\n" + f" {result}\n" + "@work.init\n" + "async def initialize():\n" + " global loop, future, gate, consumer\n" + " loop = asyncio.get_running_loop()\n" + " future = loop.create_future()\n" + " gate = asyncio.Event()\n" + " async def consume():\n" + " await gate.wait()\n" + " future.set_result(42)\n" + " consumer = asyncio.create_task(consume())\n" + ) + url = start_server() + assert wait_ready(url)[0] == 200 + status, response = request(url + endpoint, "POST", {}) + assert status == 200 + if endpoint == "/submit": + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + _, job = request(url + "/result") + if job["status"] == "DONE": + response = job["response"] + break + time.sleep(0.01) + assert response == { + "success": True, + "json_result": [42] if generator else 42, + "stdout": "", + } + + +@pytest.mark.parametrize("fails", [False, True]) +def test_idle_watchdog_waits_for_initialization(start_server, monkeypatch, fails): + entered, release, monitoring = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + original = deployed.load_handle + observed = {} + + def load(): + entered.set() + assert release.wait(5) + if fails: + raise ValueError("init failed") + return original() + + def watchdog(): + observed["contact"] = runner._last_contact["ts"] + monitoring.set() + + monkeypatch.setattr(deployed, "load_handle", load) + monkeypatch.setattr(runner, "_watchdog", watchdog) + url = start_server() + try: + assert entered.wait(5) + assert request(url + "/ping", token=None)[0] == 503 + assert request(url + "/submit", "POST", {})[0] == 503 + assert not monitoring.is_set() + finally: + finished = time.monotonic() + release.set() + assert wait_ready(url)[0] == (500 if fails else 200) + assert monitoring.wait(5) + assert observed["contact"] >= finished + + +@pytest.mark.parametrize("fails", [False, True]) +def test_runner_shutdown_closes_init_loop_and_resources(artifact, monkeypatch, fails): + import importlib + + path, manifest = artifact + module_name = manifest["resources"][0]["module"] + (path / (module_name + ".py")).write_text( + 'import asyncio\nimport runpod\napp = runpod.App("example")\n' + "cancelled = finalized = False\n" + '@app.task(name="work", cpu="cpu3c")\n' + "async def work():\n return 42\n" + "async def consume():\n" + " global cancelled\n" + " try:\n await asyncio.Event().wait()\n" + " finally:\n cancelled = True\n" + "async def chunks():\n" + " global finalized\n" + " try:\n yield 1\n" + " finally:\n finalized = True\n" + "@work.init\n" + "async def initialize():\n" + " global loop, consumer, stream\n" + " loop = asyncio.get_running_loop()\n" + " consumer = asyncio.create_task(consume())\n" + " stream = chunks()\n" + " await anext(stream)\n" + " await asyncio.sleep(0)\n" + + (' raise ValueError("init failed")\n' if fails else "") + ) + httpd = Mock() + monkeypatch.setattr(runner, "TOKEN", TOKEN) + monkeypatch.setattr(runner, "_watchdog", lambda: None) + monkeypatch.setattr(runner, "ThreadingHTTPServer", lambda *_: httpd) + runner.main() + module = importlib.import_module(module_name) + assert module.loop.is_closed() + assert module.consumer.cancelled() + assert module.cancelled + assert module.finalized + assert runner._ready is not fails + httpd.server_close.assert_called_once() + + +@pytest.mark.parametrize("failure", ["bootstrap", "init"]) +def test_readiness_error_details_require_authenticated_route( + start_server, monkeypatch, failure +): + detail = "private initialization diagnostic" + error = {"error_type": "BootstrapError", "error_message": detail} + if failure == "init": + monkeypatch.setattr( + deployed, "load_handle", Mock(side_effect=ValueError(detail)) + ) + url = start_server() + else: + url = start_server(error) + assert wait_ready(url) == ( + 500, + {"ready": False, "error_type": "BootstrapError"}, + ) + for token in (None, "wrong", TOKEN): + assert request(url + "/ping", token=token) == ( + 500, + {"ready": False, "error_type": "BootstrapError"}, + ) + for endpoint in ("/result", "/execute", "/submit", "/timeout"): + method = "GET" if endpoint == "/result" else "POST" + body = None if method == "GET" else {} + for token in (None, "wrong"): + assert request(url + endpoint, method, body, token=token) == ( + 401, + {"error": "unauthorized"}, + ) + status, response = request(url + endpoint, method, body) + assert status == 500 + assert response["ready"] is False + assert response["error_type"] == "BootstrapError" + assert detail in response["error_message"] + + +@pytest.mark.parametrize("fails_late", [False, True]) +def test_hung_init_is_reaped_and_cannot_become_ready( + start_server, monkeypatch, fails_late +): + entered, release, deleted = threading.Event(), threading.Event(), threading.Event() + initializers = [] + calls = [] + + def load(): + initializers.append(threading.current_thread()) + entered.set() + assert release.wait(5) + if fails_late: + raise ValueError("late initialization failure") + return Mock() + + def terminate(): + calls.append(time.monotonic()) + if len(calls) == 1: + return False + deleted.set() + return True + + idle_watchdog = Mock() + monkeypatch.setenv("RUNPOD_TASK_INIT_TIMEOUT", "0.15") + monkeypatch.setattr(deployed, "load_handle", load) + monkeypatch.setattr(runner, "_terminate_self", terminate) + monkeypatch.setattr(runner, "_watchdog", idle_watchdog) + monkeypatch.setattr(runner, "WATCHDOG_INTERVAL", 0.01) + url = start_server() + try: + assert entered.wait(5) + assert request(url + "/submit", "POST", {})[0] == 503 + assert wait_ready(url) == ( + 500, + {"ready": False, "error_type": "BootstrapError"}, + ) + assert deleted.wait(5) + assert len(calls) == 2 + status, error = request(url + "/result") + assert status == 500 + assert error["error_message"] == "task initialization timed out" + assert runner._execution_deadline is None + finally: + release.set() + for thread in initializers: + thread.join(5) + assert not thread.is_alive() + assert runner._ready is False + assert runner._deployed_handle is None + assert request(url + "/result") == (status, error) + idle_watchdog.assert_not_called() + + +@pytest.mark.parametrize("fails", [False, True]) +@pytest.mark.parametrize("timeout", ["0.15", "1e300"]) +def test_completed_init_disarms_initialization_timeout( + start_server, monkeypatch, fails, timeout +): + terminated = threading.Event() + monkeypatch.setenv("RUNPOD_TASK_INIT_TIMEOUT", timeout) + monkeypatch.setattr(runner, "_terminate_self", lambda: terminated.set() or True) + monkeypatch.setattr( + deployed, + "load_handle", + Mock(side_effect=ValueError("init failed")) + if fails + else Mock(return_value=Mock()), + ) + url = start_server() + assert wait_ready(url)[0] == (500 if fails else 200) + assert not terminated.wait(0.3) + + +@pytest.mark.parametrize("value", ["", "0", "-1", "nan", "inf", "text"]) +def test_invalid_initialization_timeout_fails_before_loading( + start_server, monkeypatch, value +): + load = Mock() + monkeypatch.setenv("RUNPOD_TASK_INIT_TIMEOUT", value) + monkeypatch.setattr(deployed, "load_handle", load) + url = start_server() + assert wait_ready(url)[0] == 500 + status, error = request(url + "/result") + assert status == 500 + assert ( + "RUNPOD_TASK_INIT_TIMEOUT must be finite and positive" in error["error_message"] + ) + load.assert_not_called() + + +def test_initialization_timeout_defaults_independently_of_execution_cap( + start_server, monkeypatch +): + configured = Mock(wraps=runner._configured_timeout) + monkeypatch.delenv("RUNPOD_TASK_INIT_TIMEOUT", raising=False) + monkeypatch.setenv("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", "1") + monkeypatch.setattr(runner, "_configured_timeout", configured) + assert wait_ready(start_server())[0] == 200 + configured.assert_any_call("RUNPOD_TASK_INIT_TIMEOUT", 600.0) + assert runner._timeout_cap == 1 + assert runner._execution_deadline is None + + +def test_init_completion_after_deadline_waits_for_deletion(start_server, monkeypatch): + from types import SimpleNamespace + + clock = Mock(return_value=0.0) + deleted = threading.Event() + + def load(): + clock.return_value = 1.0 + return Mock() + + monkeypatch.setenv("RUNPOD_TASK_INIT_TIMEOUT", "0.05") + monkeypatch.setattr( + runner, "time", SimpleNamespace(monotonic=clock, sleep=time.sleep) + ) + monkeypatch.setattr(deployed, "load_handle", load) + monkeypatch.setattr(runner, "_terminate_self", lambda: deleted.set() or True) + url = start_server() + assert wait_ready(url)[0] == 500 + assert deleted.wait(5) + assert runner._ready is False + assert runner._deployed_handle is None