From 53e863a718fb8d5d7c2d0f0301cf9a5c56396421 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:29:43 -0700 Subject: [PATCH 1/4] feat: execute registered flash tasks --- README.md | 39 +- src/runpod_sdk_runtime/bootstrap.py | 81 +++- src/runpod_sdk_runtime/executor.py | 18 +- src/runpod_sdk_runtime/task/deployed.py | 76 ++++ src/runpod_sdk_runtime/task/runner.py | 135 +++++- tests/test_deployed_task.py | 568 ++++++++++++++++++++++++ 6 files changed, 879 insertions(+), 38 deletions(-) create mode 100644 src/runpod_sdk_runtime/task/deployed.py create mode 100644 tests/test_deployed_task.py diff --git a/README.md b/README.md index d1e58da..942d900 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,28 @@ 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_FLASH_PROTOCOL=1`, `RUNPOD_FLASH_TASK=1`, and `RUNPOD_RUNTIME_KIND=task`. +`RUNPOD_BUILD_ID`, `FLASH_APP_NAME`, `FLASH_ENVIRONMENT`, and `FLASH_RESOURCE_NAME` +identify the selected build and resource. `FLASH_BUILD_ARTIFACT_PATH` points to its +build-specific archive. The bootstrap extracts that archive 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. Bootstrap and +initialization failures return a structured `BootstrapError`; 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 +70,21 @@ 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_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..987ed51 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,19 +91,66 @@ def _resource_name(): ) +def _validate_task_environment(): + required = { + "RUNPOD_FLASH_PROTOCOL": "1", + "RUNPOD_FLASH_TASK": "1", + "RUNPOD_RUNTIME_KIND": "task", + } + for key, value in required.items(): + if os.environ.get(key) != value: + raise PhaseError("identity", f"{key} must be {value}") + for key in ( + "FLASH_APP_NAME", + "FLASH_ENVIRONMENT", + "FLASH_RESOURCE_NAME", + "RUNPOD_BUILD_ID", + "FLASH_BUILD_ARTIFACT_PATH", + ): + 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 def _locate(): """return the app tree root, unpacking the artifact if needed.""" - if PREBUILT_APP_DIR and os.path.isdir(PREBUILT_APP_DIR): + task = os.environ.get("RUNPOD_FLASH_TASK") == "1" + if not task and PREBUILT_APP_DIR and os.path.isdir(PREBUILT_APP_DIR): _log(f"using pre-unpacked app tree at {PREBUILT_APP_DIR}") return PREBUILT_APP_DIR marker = os.path.join(APP_DIR, UNPACK_MARKER) if os.path.isfile(marker): - _log(f"app tree already unpacked at {APP_DIR}") - return APP_DIR + with open(marker) as f: + identity = f.read() + if not task or identity == ARTIFACT_PATH: + _log(f"app tree already unpacked at {APP_DIR}") + return APP_DIR + raise PhaseError("locate", "unpacked app tree belongs to another artifact") deadline = time.monotonic() + ARTIFACT_WAIT_SECONDS while not os.path.isfile(ARTIFACT_PATH): @@ -146,7 +197,7 @@ def _locate(): raise PhaseError("locate", f"failed to extract artifact: {exc}") with open(marker, "w") as f: - f.write(str(time.time())) + f.write(ARTIFACT_PATH if task else str(time.time())) _log(f"artifact extracted to {APP_DIR} in {time.monotonic() - started:.1f}s") return APP_DIR @@ -175,8 +226,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 +450,11 @@ def log_message(self, *args): def _report_error(error): _log(f"FAILED {error}") - if _runtime_kind() == "api": + if _runtime_kind() == "task" or os.environ.get("RUNPOD_FLASH_TASK") == "1": + 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 +465,15 @@ def _report_error(error): def main(): try: + task = _runtime_kind() == "task" or os.environ.get("RUNPOD_FLASH_TASK") == "1" + 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 +489,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..4cb417e 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,15 @@ def serialize_chunk(chunk, request): return response -def execute_request(request): +def execute_request(request, *, prepared=None): """run one FunctionRequest to completion, returning a response dict.""" stdout_io = io.StringIO() + strict_json = prepared is not None 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 @@ -205,7 +207,9 @@ def execute_request(request): 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..bab9b40 --- /dev/null +++ b/src/runpod_sdk_runtime/task/deployed.py @@ -0,0 +1,76 @@ +"""registered task loading and json-only execution.""" + +import importlib +import json +import os +import sys + +from runpod_sdk_runtime import bootstrap +from runpod_sdk_runtime.executor import _run_awaitable, execute_request + + +def is_deployed(): + return ( + os.environ.get("RUNPOD_FLASH_TASK") == "1" + or os.environ.get("RUNPOD_RUNTIME_KIND") == "task" + or os.environ.get("RUNPOD_FLASH_PROTOCOL") == "1" + ) + + +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: + import inspect + + result = handle._init_fn() + if inspect.isawaitable(result): + _run_awaitable(result) + 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", {})), + ) diff --git a/src/runpod_sdk_runtime/task/runner.py b/src/runpod_sdk_runtime/task/runner.py index efe9cc8..7205947 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,32 @@ def _set_timeout(timeout): _watchdog_wakeup.set() +def _start_execution(timeout): + _set_timeout(_timeout_cap) + _set_timeout(timeout) + + +def _configured_timeout(): + value = os.environ.get("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS") + if value is None: + return None + try: + timeout = float(value) + except ValueError: + timeout = float("nan") + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError( + "RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS 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 +116,14 @@ def _terminate_self(): ) return False api_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") + task = os.environ.get("RUNPOD_FLASH_TASK") == "1" + 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 +136,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 +148,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 +228,30 @@ 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): + if _bootstrap_error is not None: + self._send(500, {"ready": False, **_bootstrap_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(): + 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 +270,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 +305,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,14 +351,39 @@ 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() + _bootstrap_error = bootstrap_error + _ready = False + _deployed_handle = None + _timeout_cap = None + _production = deployed.is_deployed() + + def initialize(): + global _ready, _bootstrap_error, _deployed_handle, _timeout_cap + try: + _timeout_cap = _configured_timeout() + configure_mounts() + _deployed_handle = deployed.load_handle() + _ready = True + except Exception as exc: + _bootstrap_error = { + "error_type": "BootstrapError", + "error_message": f"task initialization failed: {exc}", + } + + if not _production and bootstrap_error is None: + configure_mounts() + _ready = True _touch_contact() threading.Thread(target=_watchdog, daemon=True).start() server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler) + if _production and bootstrap_error is None: + threading.Thread(target=initialize, daemon=True).start() server.serve_forever() diff --git a/tests/test_deployed_task.py b/tests/test_deployed_task.py new file mode 100644 index 0000000..ad864d7 --- /dev/null +++ b/tests/test_deployed_task.py @@ -0,0 +1,568 @@ +"""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_FLASH_PROTOCOL": "1", + "RUNPOD_FLASH_TASK": "1", + "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 "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_FLASH_PROTOCOL", "0"), + ("RUNPOD_FLASH_TASK", "0"), + ("RUNPOD_RUNTIME_KIND", "queue"), + ("FLASH_APP_NAME", ""), + ("FLASH_RESOURCE_NAME", ""), + ("FLASH_ENVIRONMENT", ""), + ("RUNPOD_BUILD_ID", ""), + ("FLASH_BUILD_ARTIFACT_PATH", ""), + ], +) +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 + for key in ("RUNPOD_FLASH_TASK", "RUNPOD_FLASH_PROTOCOL", "RUNPOD_RUNTIME_KIND"): + monkeypatch.delenv(key, raising=False) + 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 "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_FLASH_TASK", "RUNPOD_FLASH_PROTOCOL", "RUNPOD_RUNTIME_KIND"): + 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_FLASH_TASK", "1") + 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_FLASH_TASK", "1") + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "queue") + 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_requires_exact_archive(artifact, monkeypatch): + path, _ = artifact + prebuilt = path / "prebuilt" + prebuilt.mkdir() + monkeypatch.setattr(bootstrap, "PREBUILT_APP_DIR", str(prebuilt)) + monkeypatch.setattr(bootstrap, "ARTIFACT_PATH", str(path / "build-specific.tar.gz")) + monkeypatch.setattr(bootstrap, "APP_DIR", str(path / "app")) + monkeypatch.setattr(bootstrap, "ARTIFACT_WAIT_SECONDS", 0) + with pytest.raises(bootstrap.PhaseError, match="artifact not found"): + bootstrap._locate() + + +def test_production_bootstrap_rejects_other_artifact_marker(artifact, monkeypatch): + path, _ = artifact + (path / bootstrap.UNPACK_MARKER).write_text("other-build") + monkeypatch.setattr(bootstrap, "APP_DIR", str(path)) + with pytest.raises(bootstrap.PhaseError, match="another artifact"): + bootstrap._locate() + + +@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() From c46a9a585a514d9d4305bb0febe81d8b67ace776 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:58:47 -0700 Subject: [PATCH 2/4] fix: simplify deployed task identity (CON-1833) --- README.md | 9 +++--- src/runpod_sdk_runtime/bootstrap.py | 28 +++++------------ src/runpod_sdk_runtime/task/deployed.py | 8 ----- src/runpod_sdk_runtime/task/runner.py | 4 +-- tests/test_deployed_task.py | 40 +++++++++++++------------ 5 files changed, 36 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 942d900..4422e4c 100644 --- a/README.md +++ b/README.md @@ -42,11 +42,12 @@ shared bootstrap. Source-mode tasks start through `runpod_sdk_runtime.task.runne ## Deployed tasks The control plane starts production task pods through the shared bootstrap with -`RUNPOD_FLASH_PROTOCOL=1`, `RUNPOD_FLASH_TASK=1`, and `RUNPOD_RUNTIME_KIND=task`. +`RUNPOD_RUNTIME_KIND=task`. `RUNPOD_BUILD_ID`, `FLASH_APP_NAME`, `FLASH_ENVIRONMENT`, and `FLASH_RESOURCE_NAME` -identify the selected build and resource. `FLASH_BUILD_ARTIFACT_PATH` points to its -build-specific archive. The bootstrap extracts that archive and attaches its -vendored environment. The runner validates the manifest app, resource name, and +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. Bootstrap and diff --git a/src/runpod_sdk_runtime/bootstrap.py b/src/runpod_sdk_runtime/bootstrap.py index 987ed51..859f583 100644 --- a/src/runpod_sdk_runtime/bootstrap.py +++ b/src/runpod_sdk_runtime/bootstrap.py @@ -92,20 +92,13 @@ def _resource_name(): def _validate_task_environment(): - required = { - "RUNPOD_FLASH_PROTOCOL": "1", - "RUNPOD_FLASH_TASK": "1", - "RUNPOD_RUNTIME_KIND": "task", - } - for key, value in required.items(): - if os.environ.get(key) != value: - raise PhaseError("identity", f"{key} must be {value}") + 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", - "FLASH_BUILD_ARTIFACT_PATH", ): if not os.environ.get(key, "").strip(): raise PhaseError("identity", f"{key} is required for deployed tasks") @@ -138,19 +131,14 @@ def _task_resource(manifest): def _locate(): """return the app tree root, unpacking the artifact if needed.""" - task = os.environ.get("RUNPOD_FLASH_TASK") == "1" - if not task and PREBUILT_APP_DIR and os.path.isdir(PREBUILT_APP_DIR): + if PREBUILT_APP_DIR and os.path.isdir(PREBUILT_APP_DIR): _log(f"using pre-unpacked app tree at {PREBUILT_APP_DIR}") return PREBUILT_APP_DIR marker = os.path.join(APP_DIR, UNPACK_MARKER) if os.path.isfile(marker): - with open(marker) as f: - identity = f.read() - if not task or identity == ARTIFACT_PATH: - _log(f"app tree already unpacked at {APP_DIR}") - return APP_DIR - raise PhaseError("locate", "unpacked app tree belongs to another artifact") + _log(f"app tree already unpacked at {APP_DIR}") + return APP_DIR deadline = time.monotonic() + ARTIFACT_WAIT_SECONDS while not os.path.isfile(ARTIFACT_PATH): @@ -197,7 +185,7 @@ def _locate(): raise PhaseError("locate", f"failed to extract artifact: {exc}") with open(marker, "w") as f: - f.write(ARTIFACT_PATH if task else str(time.time())) + f.write(str(time.time())) _log(f"artifact extracted to {APP_DIR} in {time.monotonic() - started:.1f}s") return APP_DIR @@ -450,7 +438,7 @@ def log_message(self, *args): def _report_error(error): _log(f"FAILED {error}") - if _runtime_kind() == "task" or os.environ.get("RUNPOD_FLASH_TASK") == "1": + if _runtime_kind() == "task": from runpod_sdk_runtime.task.runner import main as task_main task_main(bootstrap_error=_error_payload(error)) @@ -465,7 +453,7 @@ def _report_error(error): def main(): try: - task = _runtime_kind() == "task" or os.environ.get("RUNPOD_FLASH_TASK") == "1" + task = _runtime_kind() == "task" if task: _validate_task_environment() if _resource_name(): diff --git a/src/runpod_sdk_runtime/task/deployed.py b/src/runpod_sdk_runtime/task/deployed.py index bab9b40..ef82903 100644 --- a/src/runpod_sdk_runtime/task/deployed.py +++ b/src/runpod_sdk_runtime/task/deployed.py @@ -9,14 +9,6 @@ from runpod_sdk_runtime.executor import _run_awaitable, execute_request -def is_deployed(): - return ( - os.environ.get("RUNPOD_FLASH_TASK") == "1" - or os.environ.get("RUNPOD_RUNTIME_KIND") == "task" - or os.environ.get("RUNPOD_FLASH_PROTOCOL") == "1" - ) - - def load_handle(): bootstrap._validate_task_environment() app_dir = os.environ.get("RUNPOD_APP_DIR", "/app") diff --git a/src/runpod_sdk_runtime/task/runner.py b/src/runpod_sdk_runtime/task/runner.py index 7205947..b8dc0ec 100644 --- a/src/runpod_sdk_runtime/task/runner.py +++ b/src/runpod_sdk_runtime/task/runner.py @@ -116,7 +116,7 @@ def _terminate_self(): ) return False api_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") - task = os.environ.get("RUNPOD_FLASH_TASK") == "1" + task = os.environ.get("RUNPOD_RUNTIME_KIND") == "task" mutation = "flashTaskTerminate" if task else "podTerminate" input_type = "FlashTaskTerminateInput" if task else "PodTerminateInput" payload = json.dumps( @@ -361,7 +361,7 @@ def main(bootstrap_error=None) -> None: _ready = False _deployed_handle = None _timeout_cap = None - _production = deployed.is_deployed() + _production = os.environ.get("RUNPOD_RUNTIME_KIND") == "task" def initialize(): global _ready, _bootstrap_error, _deployed_handle, _timeout_cap diff --git a/tests/test_deployed_task.py b/tests/test_deployed_task.py index ad864d7..2fdc1b0 100644 --- a/tests/test_deployed_task.py +++ b/tests/test_deployed_task.py @@ -50,8 +50,6 @@ async def initialize(): } (tmp_path / bootstrap.MANIFEST_NAME).write_text(json.dumps(manifest)) for key, value in { - "RUNPOD_FLASH_PROTOCOL": "1", - "RUNPOD_FLASH_TASK": "1", "RUNPOD_RUNTIME_KIND": "task", "RUNPOD_BUILD_ID": "build-1", "FLASH_APP_NAME": "example", @@ -246,14 +244,11 @@ def test_init_failure_has_task_error_surface(start_server, monkeypatch): @pytest.mark.parametrize( "key,value", [ - ("RUNPOD_FLASH_PROTOCOL", "0"), - ("RUNPOD_FLASH_TASK", "0"), ("RUNPOD_RUNTIME_KIND", "queue"), ("FLASH_APP_NAME", ""), ("FLASH_RESOURCE_NAME", ""), ("FLASH_ENVIRONMENT", ""), ("RUNPOD_BUILD_ID", ""), - ("FLASH_BUILD_ARTIFACT_PATH", ""), ], ) def test_invalid_production_identity_never_falls_back( @@ -299,8 +294,7 @@ def test_manifest_qualname_resolves_class_attribute(artifact): def test_production_mode_cannot_fall_back_after_startup(start_server, monkeypatch): url = start_server() assert wait_ready(url)[0] == 200 - for key in ("RUNPOD_FLASH_TASK", "RUNPOD_FLASH_PROTOCOL", "RUNPOD_RUNTIME_KIND"): - monkeypatch.delenv(key, raising=False) + monkeypatch.delenv("RUNPOD_RUNTIME_KIND") assert ( request( url + "/execute", @@ -362,7 +356,13 @@ def test_deployed_execution_never_resolves_source_or_pickle(artifact, monkeypatc def test_source_execution_ignores_production_policy_cap(start_server, monkeypatch): - for key in ("RUNPOD_FLASH_TASK", "RUNPOD_FLASH_PROTOCOL", "RUNPOD_RUNTIME_KIND"): + 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() @@ -385,7 +385,7 @@ def test_source_execution_ignores_production_policy_cap(start_server, monkeypatc "result,expected", [(True, True), (False, False), (None, False)] ) def test_production_self_termination_uses_pod_auth(monkeypatch, result, expected): - monkeypatch.setenv("RUNPOD_FLASH_TASK", "1") + 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()) @@ -405,8 +405,7 @@ def test_production_self_termination_uses_pod_auth(monkeypatch, result, expected def test_task_bootstrap_error_never_enters_queue_loop(monkeypatch): - monkeypatch.setenv("RUNPOD_FLASH_TASK", "1") - monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "queue") + monkeypatch.setenv("RUNPOD_RUNTIME_KIND", "task") main = Mock() queue = Mock() monkeypatch.setattr(runner, "main", main) @@ -416,24 +415,27 @@ def test_task_bootstrap_error_never_enters_queue_loop(monkeypatch): assert main.call_args.kwargs["bootstrap_error"]["error_type"] == "BootstrapError" -def test_production_bootstrap_requires_exact_archive(artifact, monkeypatch): +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 / "build-specific.tar.gz")) + 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) - with pytest.raises(bootstrap.PhaseError, match="artifact not found"): - bootstrap._locate() + assert bootstrap._locate() == str(prebuilt) -def test_production_bootstrap_rejects_other_artifact_marker(artifact, monkeypatch): +def test_production_bootstrap_reuses_unpacked_app(artifact, monkeypatch): path, _ = artifact - (path / bootstrap.UNPACK_MARKER).write_text("other-build") + (path / bootstrap.UNPACK_MARKER).write_text(str(time.time())) monkeypatch.setattr(bootstrap, "APP_DIR", str(path)) - with pytest.raises(bootstrap.PhaseError, match="another artifact"): - bootstrap._locate() + 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"}']) From aa97eed8a72cbd142c95a3b4db22d0cff0ce684f Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:52:29 -0700 Subject: [PATCH 3/4] fix: enforce deployed task lifecycle boundaries --- README.md | 8 +- src/runpod_sdk_runtime/executor.py | 7 +- src/runpod_sdk_runtime/task/deployed.py | 61 ++++++- src/runpod_sdk_runtime/task/runner.py | 42 ++++- tests/test_deployed_task.py | 215 ++++++++++++++++++++++++ 5 files changed, 313 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 4422e4c..d3a3db3 100644 --- a/README.md +++ b/README.md @@ -50,8 +50,12 @@ 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. Bootstrap and -initialization failures return a structured `BootstrapError`; task pods do not +`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). diff --git a/src/runpod_sdk_runtime/executor.py b/src/runpod_sdk_runtime/executor.py index 4cb417e..ade184c 100644 --- a/src/runpod_sdk_runtime/executor.py +++ b/src/runpod_sdk_runtime/executor.py @@ -177,10 +177,11 @@ def serialize_chunk(chunk, request): return response -def execute_request(request, *, prepared=None): +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: if prepared is None: prepared, error_response = resolve_request(request) @@ -196,13 +197,13 @@ def execute_request(request, *, prepared=None): 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( diff --git a/src/runpod_sdk_runtime/task/deployed.py b/src/runpod_sdk_runtime/task/deployed.py index ef82903..97d5108 100644 --- a/src/runpod_sdk_runtime/task/deployed.py +++ b/src/runpod_sdk_runtime/task/deployed.py @@ -1,12 +1,64 @@ """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 _run_awaitable, execute_request +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(): @@ -36,11 +88,7 @@ def load_handle(): raise RuntimeError("module must expose exactly one matching @app.task handle") handle = handles.pop() if handle._init_fn is not None: - import inspect - - result = handle._init_fn() - if inspect.isawaitable(result): - _run_awaitable(result) + _run_awaitable(_initialize(handle)) return handle @@ -65,4 +113,5 @@ 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 b8dc0ec..55e31ec 100644 --- a/src/runpod_sdk_runtime/task/runner.py +++ b/src/runpod_sdk_runtime/task/runner.py @@ -232,9 +232,10 @@ def _read_request(self): raise ValueError("invalid content length") return json.loads(self.rfile.read(length) or b"{}") - def _available(self): + def _available(self, *, public=False): if _bootstrap_error is not None: - self._send(500, {"ready": False, **_bootstrap_error}) + 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}) @@ -243,7 +244,7 @@ def _available(self): def do_GET(self): # noqa: N802 - BaseHTTPRequestHandler api if self.path == "/ping": - if self._available(): + if self._available(public=True): self._send(200, {"ready": True}) return if self.path == "/result": @@ -361,7 +362,20 @@ def main(bootstrap_error=None) -> None: _ready = False _deployed_handle = None _timeout_cap = None - _production = os.environ.get("RUNPOD_RUNTIME_KIND") == "task" + _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() def initialize(): global _ready, _bootstrap_error, _deployed_handle, _timeout_cap @@ -375,16 +389,26 @@ def initialize(): "error_type": "BootstrapError", "error_message": f"task initialization failed: {exc}", } + finally: + start_watchdog() if not _production and bootstrap_error is None: configure_mounts() _ready = True - _touch_contact() - threading.Thread(target=_watchdog, daemon=True).start() server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler) - if _production and bootstrap_error is None: - threading.Thread(target=initialize, daemon=True).start() - server.serve_forever() + initializer = None + try: + if _production and bootstrap_error is None: + 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 index 2fdc1b0..bc9a31c 100644 --- a/tests/test_deployed_task.py +++ b/tests/test_deployed_task.py @@ -235,6 +235,9 @@ def test_init_failure_has_task_error_surface(start_server, monkeypatch): 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) @@ -332,6 +335,9 @@ def test_invalid_configured_cap_fails_readiness(start_server, monkeypatch, 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"] @@ -568,3 +574,212 @@ def log_message(self, *_): 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"] From 6f5b1629d92b2d38ecb9a367e272099cc227ecb9 Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Fri, 23 Oct 2026 13:27:14 -0700 Subject: [PATCH 4/4] fix: bound deployed task initialization --- README.md | 8 ++ src/runpod_sdk_runtime/task/runner.py | 63 ++++++++++--- tests/test_deployed_task.py | 127 ++++++++++++++++++++++++++ 3 files changed, 185 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index d3a3db3..a6f6199 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,14 @@ 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. +`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 diff --git a/src/runpod_sdk_runtime/task/runner.py b/src/runpod_sdk_runtime/task/runner.py index 55e31ec..72d98e1 100644 --- a/src/runpod_sdk_runtime/task/runner.py +++ b/src/runpod_sdk_runtime/task/runner.py @@ -77,18 +77,16 @@ def _start_execution(timeout): _set_timeout(timeout) -def _configured_timeout(): - value = os.environ.get("RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS") +def _configured_timeout(key="RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS", default=None): + value = os.environ.get(key) if value is None: - return None + return default try: timeout = float(value) except ValueError: timeout = float("nan") if not math.isfinite(timeout) or timeout <= 0: - raise ValueError( - "RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS must be finite and positive" - ) + raise ValueError(f"{key} must be finite and positive") return timeout @@ -377,20 +375,48 @@ 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() - _deployed_handle = deployed.load_handle() - _ready = True + handle = deployed.load_handle() except Exception as exc: - _bootstrap_error = { + error = { "error_type": "BootstrapError", "error_message": f"task initialization failed: {exc}", } - finally: - start_watchdog() + 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() @@ -399,8 +425,19 @@ def initialize(): initializer = None try: if _production and bootstrap_error is None: - initializer = threading.Thread(target=initialize, daemon=True) - initializer.start() + 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() diff --git a/tests/test_deployed_task.py b/tests/test_deployed_task.py index bc9a31c..762f44a 100644 --- a/tests/test_deployed_task.py +++ b/tests/test_deployed_task.py @@ -783,3 +783,130 @@ def test_readiness_error_details_require_authenticated_route( 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