feat: execute registered flash tasks (CON-1833)

README.md

@@ -36,7 +36,33 @@ python -m runpod_sdk_runtime.bootstrap
 python -m runpod_sdk_runtime.task.runner
 ```
 
-`RUNPOD_RUNTIME_KIND` selects `queue` or `api` when running the shared bootstrap.
+`RUNPOD_RUNTIME_KIND` selects `queue`, `api`, or deployed `task` when running the
+shared bootstrap. Source-mode tasks start through `runpod_sdk_runtime.task.runner`.
+
+## Deployed tasks
+
+The control plane starts production task pods through the shared bootstrap with
+`RUNPOD_RUNTIME_KIND=task`.
+`RUNPOD_BUILD_ID`, `FLASH_APP_NAME`, `FLASH_ENVIRONMENT`, and `FLASH_RESOURCE_NAME`
+identify the selected build and resource. The host delivers the environment's
+active artifact at `/root/.runpod/artifact.tar.gz` on its read-only mount. The
+bootstrap waits for that archive, extracts it and attaches its vendored
+environment. The runner validates the manifest app, resource name, and
+task kind, then imports the registered function and runs its init hook.
+
+`GET /ping` reports readiness only after initialization succeeds. Pending
+initialization returns HTTP 503 with `{"ready": false}`. Bootstrap and
+initialization failures return HTTP 500 with
+`{"ready": false, "error_type": "BootstrapError"}`. Error details are available
+through authenticated `GET /result`. Clients should treat the bootstrap failure
+as terminal and retrieve its details using the task token. Task pods do not
+consume serverless queue jobs. `RUNPOD_TASK_TOKEN` authenticates `/submit`, `/result`,
+`/execute`, and `/timeout` on `RUNPOD_TASK_PORT` (8080 by default).
+
+Production execution accepts JSON `args` and `kwargs`, with an optional transport
+`timeout`. Results use `json_result`, with captured `stdout` and the shared error
+envelope. Requests cannot supply source, dependencies, mounts, or pickle
+serialization. Source-mode tasks use the shared live executor.
 
 ## Storage bindings
 
@@ -49,13 +75,29 @@ configuration fails startup.
 
 ## Task cleanup
 
-Active background jobs and inline executions can run indefinitely, including
-after the client disconnects. The watchdog checks for idle cleanup every 15
+Active background jobs and inline executions can run indefinitely when no
+execution timeout is configured, including after the client disconnects. The
+watchdog checks for idle cleanup every 15
 seconds. `RUNPOD_TASK_IDLE_TIMEOUT` (600 seconds by default) applies only to
 pods waiting for submission or holding uncollected results.
 
-Self-deletion requires `RUNPOD_API_KEY` and `RUNPOD_POD_ID`. The runtime exits
-only after GraphQL acknowledges `podTerminate` without errors. Transient
+`RUNPOD_TASK_INIT_TIMEOUT` bounds deployed task initialization, including storage
+setup, module imports, and init hooks. It defaults to 600 seconds and must be
+finite and positive. Initialization timeout produces a terminal `BootstrapError`
+and triggers pod deletion, with retries until deletion is acknowledged. Client
+contact cannot extend this deadline, and completion after timeout cannot make
+the pod ready. Idle monitoring begins after initialization succeeds or fails;
+the initialization deadline is independent of the execution timeout.
+
+`RUNPOD_TASK_EXECUTION_TIMEOUT_SECONDS` is an optional, finite positive production
+policy cap selected by the control plane. Its deadline starts at the first accepted
+execution. Request timeouts can shorten that deadline but cannot extend it. The
+watchdog enforces execution deadlines even while work is active.
+
+Self-deletion requires `RUNPOD_API_KEY` and `RUNPOD_POD_ID`. Production task pods
+call `flashTaskTerminate` with the internal pod key in `Authorization` and the pod ID
+in `X-RunPod-Pod-Id`; source-mode pods call `podTerminate`. The runtime exits
+only after GraphQL acknowledges deletion without errors. Transient
 failures receive three bounded attempts; any failed deletion leaves the runtime
 alive for the next watchdog pass. Container exit alone does not stop pod billing.
 

src/runpod_sdk_runtime/bootstrap.py

@@ -27,6 +27,8 @@
 job-take loop answers every job with a structured BootstrapError so
 the failure surfaces in job responses (queue endpoints), or an http
 server answers every request with the error (api endpoints).
+task pods serve readiness and execution errors over their token-protected
+http transport; their init hooks run before readiness succeeds.
 
 stdlib only: this file runs on the image's python before any artifact
 environment exists. runtime images bake it in, and custom images install
@@ -57,11 +59,13 @@
 WORKER_MODULES = {
     "queue": "runpod_sdk_runtime.queue.worker",
     "api": "runpod_sdk_runtime.api.server",
+    "task": "runpod_sdk_runtime.task.runner",
 }
 
 SDK_MODULES = {
     "queue": ("runpod.apps", "runpod.serverless"),
     "api": ("runpod.apps", "fastapi", "uvicorn"),
+    "task": ("runpod.apps", "runpod.apps.handles"),
 }
 
 
@@ -87,6 +91,41 @@ def _resource_name():
     )
 
 
+def _validate_task_environment():
+    if _runtime_kind() != "task":
+        raise PhaseError("identity", "RUNPOD_RUNTIME_KIND must be task")
+    for key in (
+        "FLASH_APP_NAME",
+        "FLASH_ENVIRONMENT",
+        "FLASH_RESOURCE_NAME",
+        "RUNPOD_BUILD_ID",
+    ):
+        if not os.environ.get(key, "").strip():
+            raise PhaseError("identity", f"{key} is required for deployed tasks")
+
+
+def _task_resource(manifest):
+    if not isinstance(manifest, dict):
+        raise PhaseError("attach", "manifest must be an object")
+    if manifest.get("app") != os.environ.get("FLASH_APP_NAME"):
+        raise PhaseError("attach", "manifest app does not match FLASH_APP_NAME")
+    resources = manifest.get("resources")
+    if not isinstance(resources, list) or any(
+        not isinstance(entry, dict) for entry in resources
+    ):
+        raise PhaseError("attach", "manifest resources must be an array of objects")
+    entries = [r for r in resources if r.get("name") == _resource_name()]
+    if len(entries) != 1 or entries[0].get("kind") != "task":
+        raise PhaseError("attach", "manifest must contain exactly one matching task")
+    entry = entries[0]
+    if any(
+        not isinstance(entry.get(key), str) or not entry[key].strip()
+        for key in ("module", "qualname")
+    ):
+        raise PhaseError("attach", "task module and qualname are required")
+    return entry
+
+
 # ---------------------------------------------------------------- locate
 
 
@@ -175,8 +214,11 @@ def _manifest(app_dir):
     path = os.path.join(app_dir, MANIFEST_NAME)
     if not os.path.isfile(path):
         raise PhaseError("attach", f"manifest not found at {path}")
-    with open(path) as f:
-        return json.load(f)
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except (OSError, ValueError) as exc:
+        raise PhaseError("attach", f"failed to read manifest: {exc}") from exc
 
 
 def _sdk_importable(paths):
@@ -396,7 +438,11 @@ def log_message(self, *args):
 
 def _report_error(error):
     _log(f"FAILED {error}")
-    if _runtime_kind() == "api":
+    if _runtime_kind() == "task":
+        from runpod_sdk_runtime.task.runner import main as task_main
+
+        task_main(bootstrap_error=_error_payload(error))
+    elif _runtime_kind() == "api":
         _api_error_server(error)
     else:
         _queue_error_loop(error)
@@ -407,10 +453,15 @@ def _report_error(error):
 
 def main():
     try:
+        task = _runtime_kind() == "task"
+        if task:
+            _validate_task_environment()
         if _resource_name():
             app_dir = _locate()
             paths = _attach(app_dir)
             manifest = _manifest(app_dir)
+            if task:
+                _task_resource(manifest)
             _verify_excluded(manifest)
             _install_system(manifest)
             ok, err = _sdk_importable(paths)
@@ -426,6 +477,10 @@ def main():
             _serve([])
     except PhaseError as exc:
         _report_error(exc)
+    except Exception as exc:
+        if not task:
+            raise
+        _report_error(PhaseError("bootstrap", str(exc)))
 
 
 if __name__ == "__main__":

src/runpod_sdk_runtime/executor.py

@@ -102,12 +102,12 @@ def _deserialize_args(request):
     return args, kwargs
 
 
-def _serialize_result(result, fmt):
+def _serialize_result(result, fmt, *, allow_nan=True):
     if fmt == "json":
         # json-format requests promise json on the wire; failing here
         # keeps dev behavior identical to a deployed endpoint
         try:
-            json.dumps(result)
+            json.dumps(result, allow_nan=allow_nan)
         except (TypeError, ValueError) as exc:
             raise TypeError(
                 f"return value must be json-serializable "
@@ -177,13 +177,16 @@ def serialize_chunk(chunk, request):
     return response
 
 
-def execute_request(request):
+def execute_request(request, *, prepared=None, run_awaitable=None):
     """run one FunctionRequest to completion, returning a response dict."""
     stdout_io = io.StringIO()
+    strict_json = prepared is not None
+    await_result = run_awaitable or _run_awaitable
     try:
-        prepared, error_response = resolve_request(request)
-        if error_response is not None:
-            return error_response
+        if prepared is None:
+            prepared, error_response = resolve_request(request)
+            if error_response is not None:
+                return error_response
         fn, args, kwargs = prepared
 
         # stdout is teed: captured for the job response and written
@@ -194,18 +197,20 @@ def execute_request(request):
         with redirect_stdout(_Tee(stdout_io, sys.__stdout__)):
             result = fn(*args, **kwargs)
             if hasattr(result, "__await__"):
-                result = _run_awaitable(result)
+                result = await_result(result)
             # generators aggregate so live .remote() matches the
             # deployed worker's return_aggregate_stream output
             elif inspect.isgenerator(result):
                 result = list(result)
             elif inspect.isasyncgen(result):
-                result = _run_awaitable(_drain_async_gen(result))
+                result = await_result(_drain_async_gen(result))
 
         response = {"success": True, "stdout": stdout_io.getvalue()}
         response.update(
             _serialize_result(
-                result, request.get("serialization_format", "cloudpickle")
+                result,
+                request.get("serialization_format", "cloudpickle"),
+                allow_nan=not strict_json,
             )
         )
         return response

src/runpod_sdk_runtime/task/deployed.py

@@ -0,0 +1,117 @@
+"""registered task loading and json-only execution."""
+
+import asyncio
+import importlib
+import inspect
+import json
+import os
+import sys
+import threading
+
+from runpod_sdk_runtime import bootstrap
+from runpod_sdk_runtime.executor import _await, execute_request
+
+_loop = None
+_loop_thread = None
+_loop_lock = threading.Lock()
+
+
+def _serve_loop(loop):
+    asyncio.set_event_loop(loop)
+    try:
+        loop.run_forever()
+    finally:
+        pending = asyncio.all_tasks(loop)
+        for task in pending:
+            task.cancel()
+        if pending:
+            loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.run_until_complete(loop.shutdown_default_executor())
+        loop.close()
+
+
+def _run_awaitable(awaitable):
+    global _loop, _loop_thread
+    with _loop_lock:
+        if _loop is None:
+            _loop = asyncio.new_event_loop()
+            _loop_thread = threading.Thread(
+                target=_serve_loop, args=(_loop,), daemon=True
+            )
+            _loop_thread.start()
+        future = asyncio.run_coroutine_threadsafe(_await(awaitable), _loop)
+    return future.result()
+
+
+def close():
+    global _loop, _loop_thread
+    with _loop_lock:
+        if _loop is None:
+            return
+        _loop.call_soon_threadsafe(_loop.stop)
+        _loop_thread.join()
+        _loop = None
+        _loop_thread = None
+
+
+async def _initialize(handle):
+    result = handle._init_fn()
+    if inspect.isawaitable(result):
+        await result
+
+
+def load_handle():
+    bootstrap._validate_task_environment()
+    app_dir = os.environ.get("RUNPOD_APP_DIR", "/app")
+    entry = bootstrap._task_resource(bootstrap._manifest(app_dir))
+    if app_dir not in sys.path:
+        sys.path.insert(0, app_dir)
+    module = importlib.import_module(entry["module"])
+
+    from runpod.apps.handles import FunctionHandle
+    from runpod.apps.spec import ResourceKind
+
+    target = module
+    for name in entry["qualname"].split("."):
+        target = getattr(target, name, None)
+    handles = {
+        value
+        for value in [target, *vars(module).values()]
+        if isinstance(value, FunctionHandle)
+        and value.spec.name == entry["name"]
+        and value.spec.kind == ResourceKind.TASK
+        and value._app.name == os.environ["FLASH_APP_NAME"]
+        and value._fn.__qualname__ == entry["qualname"]
+    }
+    if len(handles) != 1:
+        raise RuntimeError("module must expose exactly one matching @app.task handle")
+    handle = handles.pop()
+    if handle._init_fn is not None:
+        _run_awaitable(_initialize(handle))
+    return handle
+
+
+def validate_request(request, *, timeout_only=False):
+    allowed = (
+        {"timeout"}
+        if timeout_only
+        else {"args", "kwargs", "serialization_format", "timeout"}
+    )
+    if request.keys() - allowed:
+        raise ValueError("deployed tasks accept only json args, kwargs and timeout")
+    if request.get("serialization_format", "json") != "json":
+        raise ValueError("deployed tasks require json serialization")
+    if not isinstance(request.get("args", []), list):
+        raise ValueError("args must be a json array")
+    if not isinstance(request.get("kwargs", {}), dict):
+        raise ValueError("kwargs must be a json object")
+    json.dumps(request, allow_nan=False)
+
+
+def execute(handle, request):
+    return execute_request(
+        {"serialization_format": "json"},
+        prepared=(handle._fn, request.get("args", []), request.get("kwargs", {})),
+        run_awaitable=_run_awaitable,
+    )