Merge pull request #3 from runpod-workers/zeke/con-1751-alpha-feedback-apptask-pods-keep-running-and-billing-after

fix: bind worker volumes and acknowledge task pod termination

Author: KAJdev

.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
 

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.

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:

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:

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)

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)

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)