fix(serverless): keep log capture and shutdown behavior off the existing path

An audit of the most popular worker templates showed this feature reaching
workers that never asked for it. Two of them install the SDK unpinned, so they
take whatever we publish, and both fill the buffer with operational chatter while
returning errors rather than raising, so the capture was pure exposure with no
benefit. Meanwhile the workers the feature was built for run their engine as a
subprocess, whose output arrives on inherited file descriptors and was never
capturable in the first place.

Make capture opt-in. RUNPOD_LOG_CAPTURE=auto, the default, installs the proxy
only when prestart hooks are registered, so a worker that has not adopted
prestart is untouched: streams unwrapped, payload unchanged. `all` and `off`
force it either way.

Stop truncating what this feature did not add. A handler's error_message and
error_traceback are reported in full again; 16KB is well under the size of a
real framework traceback, and the job-done body limit was not this change's
problem to solve. Only the logs field stays bounded.

Restore the in-flight job-take on ordinary shutdown. The platform assigns a job
before writing the take response and does not release it when the client
disconnects, so aborting a take that already carries one strands that request
until the heartbeat sweeper reclaims it, spending its single retry. Prestart
failure already had a grace period for this; every other shutdown now gets a
short one.

Also make the proxy indistinguishable from the stream it wraps: register it on
io.TextIOBase so isinstance checks still pass (subclassing would shadow
encoding with None), and capture writelines, which bypassed write().

Document what capture cannot see: child processes, log handlers constructed
before startup, and plain threads.

Author: jasonwang-runpod

docs/serverless/worker.md

@@ -94,9 +94,9 @@ when hooks are registered. It never silently skips registered startup work.
 
 ## Failure logs
 
-When a handler or a prestart hook fails, the SDK attaches the last 16 KB of
-whatever the worker wrote to `stdout`/`stderr` to the failure it returns, under
-a `logs` key:
+When a prestart hook or a handler fails, the SDK can attach the last 16 KB of
+whatever the worker wrote to `stdout`/`stderr` to the failure it reports, under a
+`logs` key:
 
 ```json
 {
@@ -107,20 +107,36 @@ a `logs` key:
 }
 ```
 
-That payload is returned to whoever called the request, so anything the worker
-prints can reach the caller. On a private endpoint the caller is you, and the
-same output is already in your worker logs. On a public or shared endpoint the
-caller is a third party. If your worker prints credentials, connection strings,
-or other data you do not want returned, either stop printing them or turn
-capture off:
+That payload is returned to whoever called the request, so capture is **not on by
+default**. `RUNPOD_LOG_CAPTURE` controls it:
 
-```bash
-RUNPOD_SKIP_LOG_CAPTURE=true
-```
+| Value | Behavior |
+|-------|----------|
+| `auto` (default) | Capture only when prestart hooks are registered. A worker that has not adopted prestart hooks is untouched: streams are not wrapped and no `logs` key is added. |
+| `all` | Always capture, including handler failures on workers with no prestart hooks. |
+| `off` | Never capture, even with prestart hooks registered. |
+
+Only the `logs` field is bounded. `error_message` and `error_traceback` are
+reported in full, exactly as they were before this feature existed.
+
+### What capture can and cannot see
+
+Capture replaces `sys.stdout`/`sys.stderr` at worker startup, so it sees `print`
+and direct stream writes made while a hook or handler runs. It does **not** see:
+
+- **Child processes.** A subprocess inherits file descriptors 1 and 2 directly,
+  below Python. If your worker runs an engine as a subprocess and wants its output
+  captured, pipe it (`subprocess.Popen(..., stdout=PIPE)`) and re-emit the lines.
+- **Log handlers created before startup.** `logging.StreamHandler` resolves
+  `sys.stderr` once, when it is constructed. A `logging.basicConfig()` at module
+  import binds the real stream before the SDK starts, so those records bypass
+  capture. Call `logging.basicConfig()` from inside a prestart hook if you want
+  its output attached.
+- **Plain threads.** Output from a thread your code starts directly is not
+  captured; `asyncio.to_thread` is.
 
-With capture off, failures still report the exception type, message, and
-traceback; only the `logs` key is dropped. Output always continues to reach your
-worker logs either way.
+In all of these cases output still reaches your worker logs as usual — only the
+`logs` field is affected.
 
 ## Worker Refresh
 

runpod/serverless/modules/rp_capture.py

@@ -6,18 +6,30 @@
 real stream and a buffer in a contextvar.
 
 Captured output is attached to failure payloads, which are returned to whoever
-called the request. Set `RUNPOD_SKIP_LOG_CAPTURE=true` to keep worker output
-inside the worker.
+called the request, so capture is not on by default. `RUNPOD_LOG_CAPTURE`
+selects when it runs:
+
+    auto (default)  capture only when prestart hooks are registered, so a worker
+                    that has not adopted prestart behaves exactly as before
+    all             always capture, including handler failures on workers with
+                    no prestart hooks
+    off             never capture
 """
 
 import contextlib
 import contextvars
+import io
 import os
 import sys
 from collections.abc import Generator
 
 MAX_CAPTURED_CHARS = 16 * 1024
 
+CAPTURE_AUTO = "auto"
+CAPTURE_ALL = "all"
+CAPTURE_OFF = "off"
+_CAPTURE_MODES = (CAPTURE_AUTO, CAPTURE_ALL, CAPTURE_OFF)
+
 # Capture buffer for the current context
 _current: "contextvars.ContextVar[_RingBuffer | None]" = contextvars.ContextVar(
     "rp_stdio_capture", default=None
@@ -54,6 +66,11 @@ def write(self, text) -> int:
                 buffer.write(text)
         return n
 
+    def writelines(self, lines) -> None:
+        # Not covered by write(); callers that use it would otherwise bypass capture.
+        for line in lines:
+            self.write(line)
+
     def flush(self) -> None:
         self._real.flush()
 
@@ -62,15 +79,36 @@ def __getattr__(self, name):
         return getattr(self._real, name)
 
 
+# Registered rather than subclassed: inheriting io.TextIOBase would shadow
+# encoding/errors/newlines with None instead of delegating to the real stream.
+# Registration is what keeps `isinstance(sys.stdout, io.TextIOBase)` true for
+# libraries that type-check the stream.
+io.TextIOBase.register(_TeeProxy)
+
+
+def capture_mode() -> str:
+    """Resolve `RUNPOD_LOG_CAPTURE`, falling back to `auto` on anything unknown."""
+    mode = os.environ.get("RUNPOD_LOG_CAPTURE", CAPTURE_AUTO).strip().lower()
+    return mode if mode in _CAPTURE_MODES else CAPTURE_AUTO
+
+
 def install() -> None:
-    """Install the tee proxy on stdout/stderr. Idempotent.
+    """Install the tee proxy on stdout/stderr, if this worker wants capture.
 
-    Skipped when `RUNPOD_SKIP_LOG_CAPTURE=true`. Without the proxy nothing ever
-    reaches a capture buffer, so no failure payload carries worker output.
+    Idempotent. Without the proxy nothing ever reaches a capture buffer, so no
+    failure payload carries worker output and the streams are left untouched.
     """
-    if os.environ.get("RUNPOD_SKIP_LOG_CAPTURE", "").lower() == "true":
+    mode = capture_mode()
+    if mode == CAPTURE_OFF:
         return
 
+    if mode == CAPTURE_AUTO:
+        # Local import: rp_prestart imports this module.
+        from .rp_prestart import has_prestart_hooks
+
+        if not has_prestart_hooks():
+            return
+
     if not isinstance(sys.stdout, _TeeProxy):
         sys.stdout = _TeeProxy(sys.stdout)
     if not isinstance(sys.stderr, _TeeProxy):

runpod/serverless/modules/rp_job.py

@@ -15,7 +15,7 @@
 
 from ...version import __version__ as runpod_version
 from ..utils import rp_debugger
-from .rp_capture import capture, clip
+from .rp_capture import capture
 from .rp_handler import is_generator
 from .rp_http import send_result, stream_result
 from .rp_tips import check_return_size
@@ -290,8 +290,8 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]:
             captured_logs = cap.getvalue()
             error_info = {
                 "error_type": str(type(err)),
-                "error_message": clip(str(err)),
-                "error_traceback": clip(traceback.format_exc()),
+                "error_message": str(err),
+                "error_traceback": traceback.format_exc(),
                 "hostname": os.environ.get("RUNPOD_POD_HOSTNAME", "unknown"),
                 "worker_id": os.environ.get("RUNPOD_POD_ID", "unknown"),
                 "runpod_version": runpod_version,
@@ -339,6 +339,6 @@ async def run_job_generator(
             error = f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"
             if captured_logs:
                 error += f"\nlogs:\n{captured_logs}"
-            yield {"error": clip(error)}
+            yield {"error": error}
         finally:
             log.info("Finished running generator.", job["id"])

runpod/serverless/modules/rp_scale.py

@@ -70,6 +70,9 @@ def __init__(self, config: dict[str, Any]):
         self.jobs_fetcher = get_job
         self.jobs_fetcher_timeout = 90
         self.prestart_claim_timeout = 10
+        # Only needs to outlast a response already in transit, so it is far
+        # shorter than prestart_claim_timeout, which waits for a job to appear.
+        self.shutdown_take_grace = 2
         self.jobs_handler = handle_job
 
         if concurrency_modifier := config.get("concurrency_modifier"):
@@ -218,8 +221,15 @@ def current_occupancy(self) -> int:
         return current_progress_count + current_queue_count
 
     async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: int):
-        """After prestart failure, briefly keep an in-flight job-take alive so
-        its request can receive the error."""
+        """Run one blocking job-take without letting it delay shutdown.
+
+        Shutdown cannot simply abandon the take. The platform assigns a job
+        before it writes the response and does not release it when the client
+        disconnects, so aborting a take that already carries a job strands that
+        request until the worker-heartbeat sweeper reclaims it, which also
+        spends the job's one retry. Give the response a bounded window to land:
+        long enough to report a prestart failure against, short otherwise.
+        """
         fetch_task = asyncio.create_task(self.jobs_fetcher(session, jobs_needed))
         shutdown_task = asyncio.create_task(self._shutdown_event.wait())
         try:
@@ -232,14 +242,13 @@ async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: i
                 return await fetch_task
 
             if shutdown_task in done:
-                if self._prestart_error is None:
-                    return None
-
-                self._failure_take_done = True
+                if self._prestart_error is not None:
+                    self._failure_take_done = True
+                    grace = self.prestart_claim_timeout
+                else:
+                    grace = self.shutdown_take_grace
                 try:
-                    return await asyncio.wait_for(
-                        fetch_task, timeout=self.prestart_claim_timeout
-                    )
+                    return await asyncio.wait_for(fetch_task, timeout=grace)
                 except asyncio.TimeoutError:
                     return None
 

tests/test_serverless/test_capture.py

@@ -13,6 +13,10 @@
 
 from runpod.serverless.modules import rp_capture
 from runpod.serverless.modules.rp_job import run_job, run_job_generator
+from runpod.serverless.modules.rp_prestart import (
+    clear_prestart_hooks,
+    register_prestart_hook,
+)
 
 
 def _run(coro):
@@ -65,27 +69,57 @@ def test_install_is_idempotent(self):
             assert sys.stderr is installed_stderr
 
 
-class TestCaptureOptOut(unittest.TestCase):
-    """`RUNPOD_SKIP_LOG_CAPTURE` keeps worker output out of failure payloads."""
+class TestCaptureIsOptIn(unittest.TestCase):
+    """A worker that has not adopted prestart hooks must see no change at all:
+    streams untouched and no logs field in the failure it reports."""
 
-    def test_install_is_skipped(self):
-        real_stdout, real_stderr = sys.stdout, sys.stderr
+    def setUp(self):
+        clear_prestart_hooks()
+        self.addCleanup(clear_prestart_hooks)
+
+    def _install_with(self, env) -> bool:
+        """Install under `env` and report whether the proxy was applied."""
+        real_stdout, real_stderr = sys.__stdout__, sys.__stderr__
         with (
-            patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}),
+            patch.dict(os.environ, env, clear=False),
             patch.object(sys, "stdout", real_stdout),
             patch.object(sys, "stderr", real_stderr),
         ):
             rp_capture.install()
+            return isinstance(sys.stdout, rp_capture._TeeProxy) and isinstance(
+                sys.stderr, rp_capture._TeeProxy
+            )
+
+    def test_auto_skips_install_without_hooks(self):
+        assert self._install_with({"RUNPOD_LOG_CAPTURE": "auto"}) is False
+
+    def test_auto_is_the_default(self):
+        env = {k: v for k, v in os.environ.items() if k != "RUNPOD_LOG_CAPTURE"}
+        with patch.dict(os.environ, env, clear=True):
+            assert rp_capture.capture_mode() == rp_capture.CAPTURE_AUTO
+        assert self._install_with({}) is False
+
+    def test_auto_installs_once_a_hook_is_registered(self):
+        register_prestart_hook(lambda: None)
+        assert self._install_with({"RUNPOD_LOG_CAPTURE": "auto"}) is True
+
+    def test_all_installs_without_hooks(self):
+        assert self._install_with({"RUNPOD_LOG_CAPTURE": "all"}) is True
+
+    def test_off_skips_install_even_with_hooks(self):
+        register_prestart_hook(lambda: None)
+        assert self._install_with({"RUNPOD_LOG_CAPTURE": "off"}) is False
 
-            assert not isinstance(sys.stdout, rp_capture._TeeProxy)
-            assert not isinstance(sys.stderr, rp_capture._TeeProxy)
+    def test_unknown_mode_falls_back_to_auto(self):
+        with patch.dict(os.environ, {"RUNPOD_LOG_CAPTURE": "yes-please"}):
+            assert rp_capture.capture_mode() == rp_capture.CAPTURE_AUTO
 
-    def test_error_omits_logs_without_the_proxy(self):
+    def test_error_payload_is_unchanged_without_capture(self):
         def handler(_job):
             print("HF_TOKEN=hf_secret")
             raise RuntimeError("kernel panic")
 
-        with patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}):
+        with patch.dict(os.environ, {"RUNPOD_LOG_CAPTURE": "off"}):
             rp_capture.install()
             result = _run(run_job(handler, {"id": "j3"}))
 
@@ -141,8 +175,8 @@ def test_handler_success_has_no_error(self):
 
 
 class TestErrorFieldBounding(unittest.TestCase):
-    """Error strings shipped back to the platform are bounded so a huge message/log can't
-    blow past the job-done body limit."""
+    """Only the log tail this SDK adds is bounded. A handler's own error message and
+    traceback are reported exactly as before, so existing workers see no truncation."""
 
     def test_clip_keeps_head_and_tail(self):
         text = "A" * 100 + "B" * 100
@@ -155,17 +189,34 @@ def test_clip_keeps_head_and_tail(self):
     def test_clip_passthrough_when_small(self):
         assert rp_capture.clip("short", limit=100) == "short"
 
-    def test_run_job_bounds_error_message(self):
+    def test_run_job_reports_a_huge_handler_error_untruncated(self):
         huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3)
 
         def handler(_job):
             raise ValueError(huge)
 
         result = _run(run_job(handler, {"id": "big"}))
         error = json.loads(result["error"])
-        assert len(error["error_message"]) <= rp_capture.MAX_CAPTURED_CHARS + 100
+        assert error["error_message"] == huge
+        assert "truncated" not in error["error_traceback"]
 
-    def test_run_job_generator_bounds_combined_error(self):
+    def test_run_job_bounds_only_the_captured_logs(self):
+        def handler(_job):
+            print("y" * (rp_capture.MAX_CAPTURED_CHARS * 3))
+            raise ValueError("boom")
+
+        real = io.StringIO()
+        with (
+            patch.object(sys, "stdout", rp_capture._TeeProxy(real)),
+            patch.object(sys, "stderr", rp_capture._TeeProxy(real)),
+        ):
+            result = _run(run_job(handler, {"id": "big-logs"}))
+
+        error = json.loads(result["error"])
+        assert error["error_message"] == "boom"
+        assert len(error["logs"]) <= rp_capture.MAX_CAPTURED_CHARS
+
+    def test_run_job_generator_bounds_only_the_captured_logs(self):
         huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3)
 
         def handler(_job):
@@ -180,7 +231,23 @@ def handler(_job):
         ):
             result = _run_gen(run_job_generator(handler, {"id": "big-generator"}))
 
-        assert len(result[0]["error"]) <= rp_capture.MAX_CAPTURED_CHARS + 100
+        # The handler's own error survives whole; only the appended log tail is capped.
+        assert huge in result[0]["error"]
+        logs = result[0]["error"].split("\nlogs:\n", 1)[1]
+        assert len(logs) <= rp_capture.MAX_CAPTURED_CHARS
+
+    def test_proxy_is_indistinguishable_from_a_text_stream(self):
+        proxy = rp_capture._TeeProxy(io.StringIO())
+        assert isinstance(proxy, io.TextIOBase)
+        assert isinstance(proxy, io.IOBase)
+
+    def test_writelines_is_captured(self):
+        real = io.StringIO()
+        proxy = rp_capture._TeeProxy(real)
+        with patch.object(sys, "stdout", proxy), rp_capture.capture() as buf:
+            sys.stdout.writelines(["a\n", "b\n"])
+        assert buf.getvalue() == "a\nb\n"
+        assert real.getvalue() == "a\nb\n"
 
 
 if __name__ == "__main__":