diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7d9f0bc..5d93a9c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,7 +24,7 @@ jobs: - uses: actions/checkout@v7 with: repository: runpod/runpod-python - ref: feat/apps-sdk + ref: main path: runpod-python - uses: actions/setup-python@v6 @@ -81,7 +81,7 @@ jobs: tags: runpod-sdk-runtime:test build-args: | BASE_IMAGE=python:3.12-slim - RUNPOD_PACKAGE_SPEC=https://github.com/runpod/runpod-python/archive/refs/heads/feat/apps-sdk.tar.gz + RUNPOD_PACKAGE_SPEC=https://github.com/runpod/runpod-python/archive/refs/heads/main.tar.gz cache-from: type=gha,scope=queue-test cache-to: type=gha,mode=max,scope=queue-test diff --git a/CHANGELOG.md b/CHANGELOG.md index 825c32f..7cade45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1 +1,11 @@ # Changelog + +## Unreleased + +### Features + +* Initialize resolved volume bindings before user code in queue, API, and task workers. + +### Bug Fixes + +* Keep active tasks running indefinitely; retry failed idle cleanup and exit only after acknowledged pod deletion. diff --git a/README.md b/README.md index 0906a4b..d1e58da 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,27 @@ python -m runpod_sdk_runtime.task.runner `RUNPOD_RUNTIME_KIND` selects `queue` or `api` when running the shared bootstrap. +## Storage bindings + +Apps provisioning supplies `RUNPOD_MOUNTS`, a JSON array of bindings with `kind` +(`network` or `global`), declared `reference`, resolved volume `id`, and absolute +`path`. Queue, API, and task workers install these before importing user code. +The compatible Apps SDK's `volume.path` uses this worker-level map; request +payloads cannot replace it. Missing configuration means no mounts; malformed +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 +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 +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. + ## Development Install the Apps SDK checkout first, then install this project: diff --git a/src/runpod_sdk_runtime/api/server.py b/src/runpod_sdk_runtime/api/server.py index 72143b0..1415493 100644 --- a/src/runpod_sdk_runtime/api/server.py +++ b/src/runpod_sdk_runtime/api/server.py @@ -29,6 +29,8 @@ import sys from typing import Any, Optional +from runpod_sdk_runtime.mounts import configure_mounts + log = logging.getLogger("runpod_sdk_runtime.api") APP_DIR = os.environ.get("RUNPOD_APP_DIR", "/app") @@ -293,6 +295,7 @@ def _build_live_app() -> Any: def build_app() -> Any: + configure_mounts() if _is_deployed(): handle = _load_api_handle() if handle._cls is not None: diff --git a/src/runpod_sdk_runtime/mounts.py b/src/runpod_sdk_runtime/mounts.py new file mode 100644 index 0000000..241a0e4 --- /dev/null +++ b/src/runpod_sdk_runtime/mounts.py @@ -0,0 +1,15 @@ +"""install provisioning-owned filesystem bindings at worker startup.""" + +import json +import os + + +def configure_mounts() -> None: + """initialize the sdk after bootstrap, before loading any user code.""" + from runpod.apps.volume import _configure_mounts + + try: + bindings = json.loads(os.environ.get("RUNPOD_MOUNTS", "[]")) + except json.JSONDecodeError as exc: + raise RuntimeError("RUNPOD_MOUNTS must contain a valid JSON array") from exc + _configure_mounts(bindings) diff --git a/src/runpod_sdk_runtime/queue/worker.py b/src/runpod_sdk_runtime/queue/worker.py index 2f3a884..67012c9 100644 --- a/src/runpod_sdk_runtime/queue/worker.py +++ b/src/runpod_sdk_runtime/queue/worker.py @@ -24,6 +24,8 @@ import runpod +from runpod_sdk_runtime.mounts import configure_mounts + APP_DIR = os.environ.get("RUNPOD_APP_DIR", "/app") MANIFEST_NAME = "runpod_manifest.json" @@ -173,6 +175,7 @@ def _worker_config(handler) -> dict: def main() -> None: + configure_mounts() if _is_deployed(): handle = _load_deployed_handle() _run_init(handle) diff --git a/src/runpod_sdk_runtime/task/runner.py b/src/runpod_sdk_runtime/task/runner.py index 8c1744d..efe9cc8 100644 --- a/src/runpod_sdk_runtime/task/runner.py +++ b/src/runpod_sdk_runtime/task/runner.py @@ -20,41 +20,53 @@ """ import json +import math import os import sys import threading +import time +import urllib.error +import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from runpod_sdk_runtime.executor import execute_request +from runpod_sdk_runtime.mounts import configure_mounts PORT = int(os.environ.get("RUNPOD_TASK_PORT", "8080")) TOKEN = os.environ.get("RUNPOD_TASK_TOKEN", "") -# watchdog: self-terminate when the client is clearly gone. RUNNING -# jobs are never killed (terminateAfter is the runaway backstop); -# state NONE means the client died before submitting, while DONE means -# the result sat uncollected. normal flows poll every ~2s, so these never fire -# for a live client. +# active work is exempt from idle cleanup. IDLE_TIMEOUT = float(os.environ.get("RUNPOD_TASK_IDLE_TIMEOUT", "600")) WATCHDOG_INTERVAL = 15.0 +TERMINATE_ATTEMPTS = 3 +TERMINATE_TIMEOUT = 5.0 # single background job slot for /submit + /result _job_lock = threading.Lock() _job_state = {"status": "NONE", "response": None} _last_contact = {"ts": None} # set at server start -# inline /execute requests in flight; the watchdog must not terminate -# the pod while one runs (long executes outlive the idle timeout) +# inline requests are protected from idle cleanup while running. _inline_executions = {"count": 0} +_execution_deadline = None +_watchdog_wakeup = threading.Event() def _touch_contact(): import time - _last_contact["ts"] = time.time() + _last_contact["ts"] = time.monotonic() + + +def _set_timeout(timeout): + global _execution_deadline + if timeout is not None: + deadline = time.monotonic() + timeout + if _execution_deadline is None or deadline < _execution_deadline: + _execution_deadline = deadline + _watchdog_wakeup.set() def _should_self_terminate(status, last_contact, now, idle_timeout, inline=0): - """the watchdog decision: kill only provably-abandoned pods.""" if status == "RUNNING" or inline > 0: return False if last_contact is None: @@ -63,57 +75,101 @@ def _should_self_terminate(status, last_contact, now, idle_timeout, inline=0): def _terminate_self(): - """terminate this pod via the injected pod-scoped api key. - - every pod carries a RUNPOD_API_KEY scoped to itself; podTerminate - with it removes the pod entirely. exiting the process is the - fallback (stops the workload; terminateAfter finishes the job). - """ - import json as _json - import urllib.request - + """exit only after the control plane acknowledges pod deletion.""" pod_id = os.environ.get("RUNPOD_POD_ID") api_key = os.environ.get("RUNPOD_API_KEY") - if pod_id and api_key: + if not pod_id or not api_key: + sys.stderr.write( + "[task-runner] cannot delete pod: missing RUNPOD_POD_ID or RUNPOD_API_KEY\n" + ) + return False + api_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") + payload = json.dumps( + { + "query": ( + "mutation podTerminate($input: PodTerminateInput!) " + "{ podTerminate(input: $input) }" + ), + "variables": {"input": {"podId": pod_id}}, + } + ).encode() + request = urllib.request.Request( + f"{api_base.rstrip('/')}/graphql", + data=payload, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + }, + ) + for attempt in range(TERMINATE_ATTEMPTS): + retryable = False try: - api_base = os.environ.get("RUNPOD_API_BASE_URL", "https://api.runpod.io") - payload = _json.dumps( - { - "query": ( - "mutation podTerminate($input: PodTerminateInput!) " - "{ podTerminate(input: $input) }" - ), - "variables": {"input": {"podId": pod_id}}, - } - ).encode() - request = urllib.request.Request( - f"{api_base}/graphql", - data=payload, - headers={ - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - }, + with urllib.request.urlopen(request, timeout=TERMINATE_TIMEOUT) as response: # noqa: S310 + body = json.load(response) + errors = body.get("errors") if isinstance(body, dict) else None + data = body.get("data") if isinstance(body, dict) else None + if ( + not errors + and isinstance(data, dict) + and "podTerminate" in data + and data["podTerminate"] is None + ): + sys.stderr.write(f"[task-runner] pod {pod_id} deletion acknowledged\n") + os._exit(0) + return True + if isinstance(errors, list) and errors: + retryable = all( + isinstance(error, dict) + and isinstance(error.get("extensions"), dict) + and error["extensions"].get("code") + in { + "INTERNAL_SERVER_ERROR", + "SERVICE_UNAVAILABLE", + "TOO_MANY_REQUESTS", + "TIMEOUT", + } + for error in errors + ) + failure = "graphql errors" if errors else "missing deletion acknowledgment" + except urllib.error.HTTPError as exc: + retryable = exc.code == 429 or 500 <= exc.code < 600 + failure = f"http {exc.code}" + exc.close() + except (urllib.error.URLError, OSError) as exc: + retryable = True + failure = str(exc) + except (ValueError, TypeError) as exc: + failure = f"invalid deletion response: {exc}" + if not retryable or attempt == TERMINATE_ATTEMPTS - 1: + sys.stderr.write( + f"[task-runner] pod {pod_id} deletion failed: {failure}; " + "keeping runtime alive for watchdog retry\n" ) - urllib.request.urlopen(request, timeout=30) # noqa: S310 - sys.stderr.write("[task-runner] self-terminated (abandoned)\n") - except Exception: # noqa: BLE001 - fall through to process exit - pass - os._exit(0) + return False + time.sleep(0.5 * (2**attempt)) + return False def _watchdog(): import time while True: - time.sleep(WATCHDOG_INTERVAL) - if _should_self_terminate( - _job_state["status"], - _last_contact["ts"], - time.time(), - IDLE_TIMEOUT, - inline=_inline_executions["count"], - ): - _terminate_self() + with _job_lock: + now = time.monotonic() + expired = _execution_deadline is not None and now >= _execution_deadline + if expired or _should_self_terminate( + _job_state["status"], + _last_contact["ts"], + now, + IDLE_TIMEOUT, + inline=_inline_executions["count"], + ): + _terminate_self() + delay = WATCHDOG_INTERVAL + if _execution_deadline is not None and _execution_deadline > now: + delay = min(delay, _execution_deadline - now) + _watchdog_wakeup.clear() + _watchdog_wakeup.wait(delay) class Handler(BaseHTTPRequestHandler): @@ -160,21 +216,47 @@ def do_POST(self): # noqa: N802 - BaseHTTPRequestHandler api if not self._authed(): self._send(401, {"error": "unauthorized"}) return + if self.path not in ("/execute", "/submit", "/timeout"): + self._send(404, {"error": "not found"}) + return + try: + request = self._read_request() + except (ValueError, TypeError): + self._send(400, {"error": "invalid request"}) + return + if not isinstance(request, dict): + self._send(400, {"error": "invalid request"}) + 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 + ): + self._send(400, {"error": "timeout must be a finite non-negative number"}) + return + if self.path == "/timeout": + with _job_lock: + _set_timeout(timeout) + self._send(200, {"timeout": timeout}) + return if self.path == "/execute": with _job_lock: + _set_timeout(timeout) _inline_executions["count"] += 1 try: - self._send(200, execute_request(self._read_request())) + self._send(200, execute_request(request)) finally: with _job_lock: _inline_executions["count"] -= 1 return if self.path == "/submit": - request = self._read_request() with _job_lock: if _job_state["status"] == "RUNNING": self._send(409, {"error": "a job is already running"}) return + _set_timeout(timeout) _job_state["status"] = "RUNNING" _job_state["response"] = None @@ -185,9 +267,11 @@ def run(): _job_state["response"] = response threading.Thread(target=run, daemon=True).start() - self._send(200, {"status": "RUNNING"}) + response = {"status": "RUNNING"} + if timeout is not None: + response["timeout"] = timeout + self._send(200, response) return - self._send(404, {"error": "not found"}) def log_message(self, format, *args): # noqa: A002 - stdlib signature # request logging is noise in container logs (dev sessions @@ -199,6 +283,7 @@ def main() -> None: 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() server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler) diff --git a/tests/test_api_server.py b/tests/test_api_server.py index 7810f1d..ebe29fc 100644 --- a/tests/test_api_server.py +++ b/tests/test_api_server.py @@ -13,18 +13,30 @@ @pytest.fixture(autouse=True) -def clean_registry(): +def clean_registry(monkeypatch): + from runpod.apps.volume import _configure_mounts + + monkeypatch.delenv("RUNPOD_MOUNTS", raising=False) + _configure_mounts([]) _clear_registry() yield _clear_registry() + _configure_mounts([]) def _write_project(tmp_path, monkeypatch, module="main_cls"): + monkeypatch.setenv( + "RUNPOD_MOUNTS", + '[{"kind":"network","reference":"data","id":"nv-1","path":"/data"}]', + ) (tmp_path / f"{module}.py").write_text( textwrap.dedent( """ import runpod from runpod import App, init, get, post + from runpod.apps.volume import NetworkVolume + volume = NetworkVolume("data") + imported = str(volume.path) app = App("api-test") @@ -33,10 +45,14 @@ class Inference: @init def setup(self): self.model = "loaded" + self.mounted = [imported, str(volume.path)] @get("/health") def health(self): - return {"model": self.model} + return { + "model": self.model, + "mounts": self.mounted + [str(volume.path)], + } @post("/generate") async def generate(self, body: dict): @@ -79,7 +95,7 @@ def test_routes_and_init(self, tmp_path, monkeypatch): assert response.json() == {"status": "healthy"} response = client.get("/health") - assert response.json() == {"model": "loaded"} + assert response.json() == {"model": "loaded", "mounts": ["/data"] * 3} response = client.post("/generate", json={"prompt": "hi"}) assert response.json() == { diff --git a/tests/test_queue_worker.py b/tests/test_queue_worker.py index b175e72..9e66f80 100644 --- a/tests/test_queue_worker.py +++ b/tests/test_queue_worker.py @@ -9,6 +9,16 @@ from runpod_sdk_runtime.queue import worker +@pytest.fixture(autouse=True) +def clean_mounts(monkeypatch): + from runpod.apps.volume import _configure_mounts + + monkeypatch.delenv("RUNPOD_MOUNTS", raising=False) + _configure_mounts([]) + yield + _configure_mounts([]) + + class TestModeDetection: def test_resource_name_flash_env(self, monkeypatch): monkeypatch.setenv("FLASH_RESOURCE_NAME", "chat") @@ -192,21 +202,20 @@ def gen_handler(job): class TestMain: - def test_live_mode(self, monkeypatch): - monkeypatch.delenv("FLASH_RESOURCE_NAME", raising=False) - monkeypatch.delenv("RUNPOD_RESOURCE_NAME", raising=False) - with patch("runpod.serverless.start") as start: - worker.main() - config = start.call_args[0][0] - assert config["handler"] is worker._live_handler - - def test_deployed_mode(self, monkeypatch, tmp_path): + async def test_deployed_mode(self, monkeypatch, tmp_path): + monkeypatch.setenv( + "RUNPOD_MOUNTS", + '[{"kind":"global","reference":"gv-1","id":"gv-1","path":"/shared"}]', + ) module_src = ( "import runpod\n" + "from runpod.apps.volume import GlobalVolume\n" + "imported = str(GlobalVolume('gv-1').path)\n" "app = runpod.App('t')\n" "@app.queue(name='greet', gpu='4090')\n" - "def greet(name: str):\n" - " return f'hi {name}'\n" + "async def greet(name: str):\n" + " yield {'name': name, 'imported': imported}\n" + " yield {'mounted': str(GlobalVolume('gv-1').path)}\n" ) (tmp_path / "user_mod_c.py").write_text(module_src) (tmp_path / worker.MANIFEST_NAME).write_text( @@ -218,25 +227,27 @@ def test_deployed_mode(self, monkeypatch, tmp_path): with patch("runpod.serverless.start") as start: worker.main() handler = start.call_args[0][0]["handler"] - assert inspect.iscoroutinefunction(handler) - - async def test_live_handler_plain_function(self): - def fn(x): - return {"ok": x} - - with ( - patch( - "runpod_sdk_runtime.executor.resolve_request", - return_value=((fn, [1], {}), None), - ), - patch( - "runpod_sdk_runtime.executor.execute_request", - return_value={"success": True, "json_result": {"ok": 1}}, - ) as execute, - ): - chunks = [c async for c in worker._live_handler({"input": {"foo": 1}})] - assert chunks == [{"success": True, "json_result": {"ok": 1}}] - execute.assert_called_once_with({"foo": 1}) + chunks = [chunk async for chunk in handler({"input": {"name": "reader"}})] + assert chunks == [ + {"name": "reader", "imported": "/shared"}, + {"mounted": "/shared"}, + ] + + async def test_live_mode(self, monkeypatch): + monkeypatch.delenv("FLASH_RESOURCE_NAME", raising=False) + monkeypatch.delenv("RUNPOD_RESOURCE_NAME", raising=False) + with patch("runpod.serverless.start") as start: + worker.main() + handler = start.call_args[0][0]["handler"] + request = { + "function_name": "echo", + "function_code": "def echo(x):\n return {'ok': x}", + "args": [1], + "serialization_format": "json", + } + chunks = [chunk async for chunk in handler({"input": request})] + assert chunks[0]["success"] is True + assert chunks[0]["json_result"] == {"ok": 1} async def test_live_handler_resolve_error(self): with patch( diff --git a/tests/test_task_runner_http.py b/tests/test_task_runner_http.py index 9a6dab9..f6dadb8 100644 --- a/tests/test_task_runner_http.py +++ b/tests/test_task_runner_http.py @@ -8,6 +8,7 @@ import base64 import json import threading +import time import urllib.error import urllib.request from http.server import ThreadingHTTPServer @@ -16,7 +17,7 @@ import pytest from runpod_sdk_runtime.task import runner as task_runner -from runpod_sdk_runtime.task.runner import Handler +from runpod_sdk_runtime.task.runner import Handler, _watchdog # http.server's per-request sockets are collected lazily; the unraisable # checker flags them as ResourceWarnings non-deterministically @@ -29,16 +30,34 @@ @pytest.fixture() def server(monkeypatch): + from runpod.apps.volume import _configure_mounts + monkeypatch.setattr(task_runner, "TOKEN", TOKEN) monkeypatch.setattr(task_runner, "_job_state", {"status": "NONE", "response": None}) + monkeypatch.setattr(task_runner, "_execution_deadline", None) + monkeypatch.setattr(task_runner, "_watchdog_wakeup", threading.Event()) + monkeypatch.setenv( + "RUNPOD_MOUNTS", + '[{"kind":"network","reference":"data","id":"nv-1","path":"/data"}]', + ) httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) httpd.daemon_threads = True - thread = threading.Thread(target=httpd.serve_forever, daemon=True) + ready = threading.Event() + + def started_server(*_): + ready.set() + return httpd + + monkeypatch.setattr(task_runner, "ThreadingHTTPServer", started_server) + monkeypatch.setattr(task_runner, "_watchdog", lambda: None) + thread = threading.Thread(target=task_runner.main, daemon=True) thread.start() + assert ready.wait(timeout=5) yield f"http://127.0.0.1:{httpd.server_port}" httpd.shutdown() thread.join(timeout=5) httpd.server_close() + _configure_mounts([]) def _request(url, method="GET", body=None, token=TOKEN): @@ -118,14 +137,158 @@ def test_submit_and_result(server): assert body["response"]["json_result"] == "done" +def test_mounts_survive_threaded_execution(server, monkeypatch): + monkeypatch.setenv("RUNPOD_MOUNTS", "[]") + status, body = _request( + f"{server}/execute", + method="POST", + body={ + "function_name": "mounted", + "function_code": ( + "from runpod.apps.volume import NetworkVolume\n" + "imported = str(NetworkVolume('data').path)\n" + "async def mounted():\n" + " yield imported\n" + " yield str(NetworkVolume('nv-1').path)\n" + ), + "serialization_format": "json", + "mounts": [], + }, + ) + assert status == 200 + assert body["success"] is True + assert body["json_result"] == ["/data", "/data"] + + def test_unknown_path_404(server): with pytest.raises(urllib.error.HTTPError) as exc_info: _request(f"{server}/nope") assert exc_info.value.code == 404 +@pytest.mark.parametrize("path", ["submit", "execute", "timeout"]) +@pytest.mark.parametrize("timeout", [-1, True, "60", float("nan"), float("inf")]) +def test_invalid_timeout_is_rejected(server, path, timeout): + with pytest.raises(urllib.error.HTTPError) as exc_info: + _request(f"{server}/{path}", method="POST", body={"timeout": timeout}) + assert exc_info.value.code == 400 + assert task_runner._execution_deadline is None + + +def test_timeout_requires_auth(server): + with pytest.raises(urllib.error.HTTPError) as exc_info: + _request(f"{server}/timeout", method="POST", body={"timeout": 0}, token="wrong") + assert exc_info.value.code == 401 + assert task_runner._execution_deadline is None + + +def test_timeout_cannot_be_extended_by_requests(server): + start = time.monotonic() + _request(f"{server}/timeout", method="POST", body={"timeout": 60}) + deadline = task_runner._execution_deadline + assert start + 60 <= deadline <= time.monotonic() + 60 + _request(f"{server}/timeout", method="POST", body={"timeout": 600}) + _request(f"{server}/timeout", method="POST", body={"timeout": None}) + _request(f"{server}/result") + assert task_runner._execution_deadline == deadline + _request(f"{server}/timeout", method="POST", body={"timeout": 0}) + assert task_runner._execution_deadline < deadline + + +@pytest.mark.parametrize("path", ["submit", "execute"]) +def test_execution_timeout_is_armed_before_user_code(server, monkeypatch, path): + deadlines = [] + entered = threading.Event() + + def execute(request): + deadlines.append(task_runner._execution_deadline) + entered.set() + return {"success": True, "json_result": 42} + + monkeypatch.setattr(task_runner, "execute_request", execute) + start = time.monotonic() + status, body = _request(f"{server}/{path}", method="POST", body={"timeout": 60}) + assert status == 200 + if path == "submit": + assert body["timeout"] == 60 + assert entered.wait(2) + assert deadlines and start + 60 <= deadlines[0] <= time.monotonic() + 60 + + +@pytest.mark.parametrize("mode", ["submit", "wait"]) +def test_timeout_survives_client_exit(server, monkeypatch, mode): + release = threading.Event() + finished = threading.Event() + terminated = threading.Event() + + def execute(request): + try: + release.wait(5) + return {"success": True, "json_result": 42} + finally: + finished.set() + + def terminate(): + terminated.set() + raise SystemExit + + def watchdog(): + try: + _watchdog() + except SystemExit: + return + + monkeypatch.setattr(task_runner, "execute_request", execute) + monkeypatch.setattr(task_runner, "_terminate_self", terminate) + thread = threading.Thread(target=watchdog, daemon=True) + try: + if mode == "submit": + _request(f"{server}/submit", method="POST", body={"timeout": 0.2}) + else: + _request(f"{server}/submit", method="POST", body={}) + _request(f"{server}/timeout", method="POST", body={"timeout": 0.2}) + assert task_runner._job_state["status"] == "RUNNING" + thread.start() + assert terminated.wait(2) + finally: + release.set() + if thread.ident is not None: + thread.join(timeout=2) + assert finished.wait(2) + + +@pytest.mark.parametrize("status,inline", [("RUNNING", 0), ("NONE", 1), ("DONE", 0)]) +def test_watchdog_enforces_deadline_despite_recent_contact(monkeypatch, status, inline): + monkeypatch.setattr(task_runner, "_execution_deadline", 10) + monkeypatch.setattr(task_runner, "_job_state", {"status": status}) + monkeypatch.setattr(task_runner, "_inline_executions", {"count": inline}) + monkeypatch.setattr(task_runner, "_last_contact", {"ts": 10}) + monkeypatch.setattr(task_runner.time, "monotonic", lambda: 10) + + def terminate(): + raise SystemExit + + monkeypatch.setattr(task_runner, "_terminate_self", terminate) + with pytest.raises(SystemExit): + _watchdog() + + +def test_expired_deadline_retries_failed_deletion(monkeypatch): + from unittest.mock import Mock + + monkeypatch.setattr(task_runner, "_execution_deadline", 0) + monkeypatch.setattr(task_runner, "_job_state", {"status": "RUNNING"}) + monkeypatch.setattr(task_runner, "_watchdog_wakeup", Mock()) + terminate = Mock(side_effect=[False, SystemExit]) + monkeypatch.setattr(task_runner, "_terminate_self", terminate) + with pytest.raises(SystemExit): + _watchdog() + assert terminate.call_count == 2 + assert task_runner._execution_deadline == 0 + + class TestWatchdog: - def test_running_job_never_killed(self): + def test_running_job_is_exempt_from_idle_cleanup(self): from runpod_sdk_runtime.task.runner import _should_self_terminate assert not _should_self_terminate("RUNNING", 0, 10_000, 600) @@ -151,56 +314,37 @@ def test_no_contact_recorded_yet(self): assert not _should_self_terminate("NONE", None, 10_000, 600) + def test_inline_execution_is_exempt_from_idle_cleanup(self): + assert not task_runner._should_self_terminate("NONE", 0, 10_000, 600, inline=1) + def test_authed_requests_touch_contact(self, server, monkeypatch): monkeypatch.setattr(task_runner, "_last_contact", {"ts": None}) _request(f"{server}/result") assert task_runner._last_contact["ts"] is not None -class TestSelfTermination: - def test_terminate_self_calls_graphql_then_exits(self, monkeypatch): - import contextlib - import io as _io - - calls = [] - - @contextlib.contextmanager - def _response(): - yield _io.BytesIO(b"{}") - - def fake_urlopen(req, timeout=None): - calls.append((req.full_url, req.data)) - return _response() - - exits = [] - monkeypatch.setenv("RUNPOD_POD_ID", "pod-1") - monkeypatch.setenv("RUNPOD_API_KEY", "pod-scoped-key") - monkeypatch.setattr("urllib.request.urlopen", fake_urlopen) - monkeypatch.setattr(task_runner.os, "_exit", lambda code: exits.append(code)) - task_runner._terminate_self() - - assert exits == [0] - url, body = calls[0] - assert url.endswith("/graphql") - assert b"podTerminate" in body - assert b"pod-1" in body - - def test_terminate_self_exits_without_credentials(self, monkeypatch): - exits = [] - monkeypatch.delenv("RUNPOD_POD_ID", raising=False) - monkeypatch.delenv("RUNPOD_API_KEY", raising=False) - monkeypatch.setattr(task_runner.os, "_exit", lambda code: exits.append(code)) - task_runner._terminate_self() - assert exits == [0] - - def test_terminate_self_survives_api_failure(self, monkeypatch): - def fail(req, timeout=None): - raise OSError("network down") - - exits = [] - monkeypatch.setenv("RUNPOD_POD_ID", "pod-1") - monkeypatch.setenv("RUNPOD_API_KEY", "k") - monkeypatch.setattr("urllib.request.urlopen", fail) - monkeypatch.setattr(task_runner.os, "_exit", lambda code: exits.append(code)) - task_runner._terminate_self() - assert exits == [0] +def test_deletion_requires_acknowledgment_and_recovers(monkeypatch): + import io + from unittest.mock import Mock + + responses = [ + io.BytesIO(b'{"errors":[{"message":"denied"}],"data":{"podTerminate":null}}'), + *[OSError("network down")] * task_runner.TERMINATE_ATTEMPTS, + urllib.error.URLError("connection reset"), + io.BytesIO(b'{"data":{"podTerminate":null}}'), + ] + exits = [] + monkeypatch.setenv("RUNPOD_POD_ID", "pod-1") + monkeypatch.setenv("RUNPOD_API_KEY", "pod-scoped-key") + monkeypatch.setattr( + task_runner.urllib.request, "urlopen", Mock(side_effect=responses) + ) + monkeypatch.setattr(task_runner.os, "_exit", exits.append) + monkeypatch.setattr(task_runner.time, "sleep", lambda _: None) + + assert task_runner._terminate_self() is False + assert not exits + assert task_runner._terminate_self() is False + assert not exits + assert task_runner._terminate_self() is True + assert exits == [0]