modal 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

modal/_container_entrypoint.py

@@ -450,7 +450,10 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
                     f"Function has {len(service.service_deps)} dependencies"
                     f" but container got {len(dep_object_ids)} object ids.\n"
                     f"Code deps: {service.service_deps}\n"
-                    f"Object ids: {dep_object_ids}"
+                    f"Object ids: {dep_object_ids}\n"
+                    "\n"
+                    "This can happen if you are defining Modal objects under a conditional statement "
+                    "that evaluates differently in the local and remote environments."
                 )
             for object_id, obj in zip(dep_object_ids, service.service_deps):
                 metadata: Message = container_app.object_handle_metadata[object_id]

modal/_partial_function.py

@@ -19,7 +19,7 @@ from ._functions import _Function
 from ._utils.async_utils import synchronizer
 from ._utils.deprecation import deprecation_warning
 from ._utils.function_utils import callable_has_non_self_params
-from .config import logger
+from .config import config, logger
 from .exception import InvalidError
 
 MAX_MAX_BATCH_SIZE = 1000
@@ -93,6 +93,26 @@ NullaryFuncOrMethod = Union[Callable[[], Any], Callable[[Any], Any]]
 NullaryMethod = Callable[[Any], Any]
 
 
+def verify_concurrent_params(params: _PartialFunctionParams, is_flash: bool = False) -> None:
+    def _verify_concurrent_params_with_flash_settings(params: _PartialFunctionParams) -> None:
+        if params.max_concurrent_inputs is not None:
+            raise TypeError(
+                "@modal.concurrent(max_inputs=...) is not yet supported for Flash functions. "
+                "Use `@modal.concurrent(target_inputs=...)` instead."
+            )
+        if params.target_concurrent_inputs is None:
+            raise TypeError("`@modal.concurrent()` missing required argument: `target_inputs`.")
+
+    def _verify_concurrent_params(params: _PartialFunctionParams) -> None:
+        if params.max_concurrent_inputs is None:
+            raise TypeError("`@modal.concurrent()` missing required argument: `max_inputs`.")
+
+    if is_flash:
+        _verify_concurrent_params_with_flash_settings(params)
+    else:
+        _verify_concurrent_params(params)
+
+
 class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
     """Object produced by a decorator in the `modal` namespace
 
@@ -378,6 +398,7 @@ def _fastapi_endpoint(
         method=method,
         web_endpoint_docs=docs,
         requested_suffix=label or "",
+        ephemeral_suffix=config.get("dev_suffix"),
         async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
         custom_domains=_parse_custom_domains(custom_domains),
         requires_proxy_auth=requires_proxy_auth,
@@ -446,6 +467,7 @@ def _web_endpoint(
         method=method,
         web_endpoint_docs=docs,
         requested_suffix=label or "",
+        ephemeral_suffix=config.get("dev_suffix"),
         async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
         custom_domains=_parse_custom_domains(custom_domains),
         requires_proxy_auth=requires_proxy_auth,
@@ -505,6 +527,7 @@ def _asgi_app(
     webhook_config = api_pb2.WebhookConfig(
         type=api_pb2.WEBHOOK_TYPE_ASGI_APP,
         requested_suffix=label or "",
+        ephemeral_suffix=config.get("dev_suffix"),
         async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
         custom_domains=_parse_custom_domains(custom_domains),
         requires_proxy_auth=requires_proxy_auth,
@@ -562,6 +585,7 @@ def _wsgi_app(
     webhook_config = api_pb2.WebhookConfig(
         type=api_pb2.WEBHOOK_TYPE_WSGI_APP,
         requested_suffix=label or "",
+        ephemeral_suffix=config.get("dev_suffix"),
         async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
         custom_domains=_parse_custom_domains(custom_domains),
         requires_proxy_auth=requires_proxy_auth,
@@ -623,6 +647,7 @@ def _web_server(
     webhook_config = api_pb2.WebhookConfig(
         type=api_pb2.WEBHOOK_TYPE_WEB_SERVER,
         requested_suffix=label or "",
+        ephemeral_suffix=config.get("dev_suffix"),
         async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
         custom_domains=_parse_custom_domains(custom_domains),
         web_server_port=port,
@@ -760,7 +785,7 @@ def _batched(
 def _concurrent(
     _warn_parentheses_missing=None,  # mdmd:line-hidden
     *,
-    max_inputs: int,  # Hard limit on each container's input concurrency
+    max_inputs: Optional[int] = None,  # Hard limit on each container's input concurrency
     target_inputs: Optional[int] = None,  # Input concurrency that Modal's autoscaler should target
 ) -> Callable[
     [Union[Callable[P, ReturnType], _PartialFunction[P, ReturnType, ReturnType]]],
@@ -812,7 +837,7 @@ def _concurrent(
             "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@modal.concurrent()`."
         )
 
-    if target_inputs and target_inputs > max_inputs:
+    if max_inputs is not None and target_inputs is not None and target_inputs > max_inputs:
         raise InvalidError("`target_inputs` parameter cannot be greater than `max_inputs`.")
 
     flags = _PartialFunctionFlags.CONCURRENT

modal/_utils/function_utils.py

@@ -75,6 +75,10 @@ def is_global_object(object_qual_name: str):
     return "<locals>" not in object_qual_name.split(".")
 
 
+def is_flash_object(experimental_options: Optional[dict[str, Any]]) -> bool:
+    return experimental_options.get("flash", False) if experimental_options else False
+
+
 def is_method_fn(object_qual_name: str):
     # methods have names like Cls.foo.
     if "<locals>" in object_qual_name:

modal/_utils/task_command_router_client.py

@@ -0,0 +1,537 @@
+# Copyright Modal Labs 2025
+import asyncio
+import base64
+import json
+import ssl
+import time
+import urllib.parse
+from typing import AsyncIterator, Optional
+
+import grpclib.client
+import grpclib.config
+import grpclib.events
+from grpclib import GRPCError, Status
+from grpclib.exceptions import StreamTerminatedError
+
+from modal.config import config, logger
+from modal.exception import ExecTimeoutError
+from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
+from modal_proto.task_command_router_grpc import TaskCommandRouterStub
+
+from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES, connect_channel, retry_transient_errors
+
+
+def _b64url_decode(data: str) -> bytes:
+    """Decode a base64url string with missing padding tolerated."""
+    padding = "=" * (-len(data) % 4)
+    return base64.urlsafe_b64decode(data + padding)
+
+
+def _parse_jwt_expiration(jwt_token: str) -> Optional[float]:
+    """Parse exp from a JWT without verification. Returns UNIX time seconds or None.
+
+    This is best-effort; if parsing fails or claim missing, returns None.
+    """
+    try:
+        parts = jwt_token.split(".")
+        if len(parts) != 3:
+            return None
+        payload_b = _b64url_decode(parts[1])
+        payload = json.loads(payload_b)
+        exp = payload.get("exp")
+        if isinstance(exp, (int, float)):
+            return float(exp)
+    except Exception:
+        # Avoid raising on malformed tokens; fall back to server-driven refresh logic.
+        logger.warning("Failed to parse JWT expiration")
+        return None
+    return None
+
+
+async def call_with_retries_on_transient_errors(
+    func,
+    *,
+    base_delay_secs: float = 0.01,
+    delay_factor: float = 2,
+    max_retries: Optional[int] = 10,
+):
+    """Call func() with transient error retries and exponential backoff.
+
+    Authentication retries are expected to be handled by the caller.
+    """
+    delay_secs = base_delay_secs
+    num_retries = 0
+
+    async def sleep_and_update_delay_and_num_retries_remaining(e: Exception):
+        nonlocal delay_secs, num_retries
+        logger.debug(f"Retrying RPC with delay {delay_secs}s due to error: {e}")
+        await asyncio.sleep(delay_secs)
+        delay_secs *= delay_factor
+        num_retries += 1
+
+    while True:
+        try:
+            return await func()
+        except GRPCError as e:
+            if (max_retries is None or num_retries < max_retries) and e.status in RETRYABLE_GRPC_STATUS_CODES:
+                await sleep_and_update_delay_and_num_retries_remaining(e)
+            else:
+                raise e
+        except AttributeError as e:
+            # StreamTerminatedError are not properly raised in grpclib<=0.4.7
+            # fixed in https://github.com/vmagamedov/grpclib/issues/185
+            # TODO: update to newer version (>=0.4.8) once stable
+            if (max_retries is None or num_retries < max_retries) and "_write_appdata" in str(e):
+                await sleep_and_update_delay_and_num_retries_remaining(e)
+            else:
+                raise e
+        except StreamTerminatedError as e:
+            if max_retries is None or num_retries < max_retries:
+                await sleep_and_update_delay_and_num_retries_remaining(e)
+            else:
+                raise e
+        except (OSError, asyncio.TimeoutError) as e:
+            if max_retries is None or num_retries < max_retries:
+                await sleep_and_update_delay_and_num_retries_remaining(e)
+            else:
+                raise ConnectionError(str(e))
+
+
+async def fetch_command_router_access(server_client, task_id: str) -> api_pb2.TaskGetCommandRouterAccessResponse:
+    """Fetch direct command router access info from Modal server."""
+    return await retry_transient_errors(
+        server_client.stub.TaskGetCommandRouterAccess,
+        api_pb2.TaskGetCommandRouterAccessRequest(task_id=task_id),
+    )
+
+
+class TaskCommandRouterClient:
+    """
+    Client used to talk directly to TaskCommandRouter service on worker hosts.
+
+    A new instance should be created per task.
+    """
+
+    @classmethod
+    async def try_init(
+        cls,
+        server_client,
+        task_id: str,
+    ) -> Optional["TaskCommandRouterClient"]:
+        """Attempt to initialize a TaskCommandRouterClient by fetching direct access.
+
+        Returns None if command router access is not enabled (FAILED_PRECONDITION).
+        """
+        try:
+            resp = await fetch_command_router_access(server_client, task_id)
+        except GRPCError as exc:
+            if exc.status == Status.FAILED_PRECONDITION:
+                logger.debug(f"Command router access is not enabled for task {task_id}")
+                return None
+            raise
+
+        logger.debug(f"Using command router access for task {task_id}")
+
+        # Build and connect a channel to the task command router now that we have access info.
+        o = urllib.parse.urlparse(resp.url)
+        if o.scheme != "https":
+            raise ValueError(f"Task router URL must be https, got: {resp.url}")
+
+        host, _, port_str = o.netloc.partition(":")
+        port = int(port_str) if port_str else 443
+        ssl_context = ssl.create_default_context()
+
+        # Allow insecure TLS when explicitly enabled via config.
+        if config["task_command_router_insecure"]:
+            logger.warning("Using insecure TLS for task command router due to MODAL_TASK_COMMAND_ROUTER_INSECURE")
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = ssl.CERT_NONE
+
+        channel = grpclib.client.Channel(
+            host,
+            port,
+            ssl=ssl_context,
+            config=grpclib.config.Configuration(
+                http2_connection_window_size=64 * 1024 * 1024,  # 64 MiB
+                http2_stream_window_size=64 * 1024 * 1024,  # 64 MiB
+            ),
+        )
+
+        await connect_channel(channel)
+
+        return cls(server_client, task_id, resp.url, resp.jwt, channel)
+
+    def __init__(
+        self,
+        server_client,
+        task_id: str,
+        server_url: str,
+        jwt: str,
+        channel: grpclib.client.Channel,
+        *,
+        stream_stdio_retry_delay_secs: float = 0.01,
+        stream_stdio_retry_delay_factor: float = 2,
+        stream_stdio_max_retries: int = 10,
+    ) -> None:
+        """Callers should not use this directly. Use TaskCommandRouterClient.try_init() instead."""
+        # Attach bearer token on all requests to the worker-side router service.
+        self._server_client = server_client
+        self._task_id = task_id
+        self._server_url = server_url
+        self._jwt = jwt
+        self._channel = channel
+        # Retry configuration for stdio streaming
+        self.stream_stdio_retry_delay_secs = stream_stdio_retry_delay_secs
+        self.stream_stdio_retry_delay_factor = stream_stdio_retry_delay_factor
+        self.stream_stdio_max_retries = stream_stdio_max_retries
+
+        # JWT refresh coordination
+        self._jwt_exp: Optional[float] = _parse_jwt_expiration(jwt)
+        self._jwt_refresh_lock = asyncio.Lock()
+        self._jwt_refresh_event = asyncio.Event()
+        self._closed = False
+
+        # Start background task to eagerly refresh JWT 30s before expiration.
+        self._jwt_refresh_task = asyncio.create_task(self._jwt_refresh_loop())
+
+        async def send_request(event: grpclib.events.SendRequest) -> None:
+            # This will get the most recent JWT for every request. No need to
+            # lock _jwt_refresh_lock: reads and writes happen on the
+            # single-threaded event loop and variable assignment is atomic.
+            event.metadata["authorization"] = f"Bearer {self._jwt}"
+
+        grpclib.events.listen(self._channel, grpclib.events.SendRequest, send_request)
+
+        self._stub = TaskCommandRouterStub(self._channel)
+
+    def __del__(self) -> None:
+        """Clean up the client when it's garbage collected."""
+        if self._closed:
+            return
+
+        self._jwt_refresh_task.cancel()
+
+        try:
+            self._channel.close()
+        except Exception:
+            pass
+
+    async def close(self) -> None:
+        """Close the client and stop the background JWT refresh task."""
+        if self._closed:
+            return
+
+        self._closed = True
+        self._jwt_refresh_task.cancel()
+        try:
+            logger.debug(f"Waiting for JWT refresh task to complete for exec with task ID {self._task_id}")
+            await self._jwt_refresh_task
+        except asyncio.CancelledError:
+            pass
+        self._channel.close()
+
+    async def exec_start(self, request: sr_pb2.TaskExecStartRequest) -> sr_pb2.TaskExecStartResponse:
+        """Start an exec'd command, properly retrying on transient errors."""
+        return await call_with_retries_on_transient_errors(
+            lambda: self._call_with_auth_retry(self._stub.TaskExecStart, request)
+        )
+
+    async def exec_stdio_read(
+        self,
+        task_id: str,
+        exec_id: str,
+        # Quotes around the type required for protobuf 3.19.
+        file_descriptor: "api_pb2.FileDescriptor.ValueType",
+        deadline: Optional[float] = None,
+    ) -> AsyncIterator[sr_pb2.TaskExecStdioReadResponse]:
+        """Stream stdout/stderr batches from the task, properly retrying on transient errors.
+
+        Args:
+            task_id: The task ID of the task running the exec'd command.
+            exec_id: The execution ID of the command to read from.
+            file_descriptor: The file descriptor to read from.
+            deadline: The deadline by which all output must be streamed. If
+              None, wait forever. If the deadline is exceeded, raises an
+              ExecTimeoutError.
+        Returns:
+            AsyncIterator[sr_pb2.TaskExecStdioReadResponse]: A stream of stdout/stderr batches.
+        Raises:
+            ExecTimeoutError: If the deadline is exceeded.
+            Other errors: If retries are exhausted on transient errors or if there's an error
+              from the RPC itself.
+        """
+        if file_descriptor == api_pb2.FILE_DESCRIPTOR_STDOUT:
+            sr_fd = sr_pb2.TASK_EXEC_STDIO_FILE_DESCRIPTOR_STDOUT
+        elif file_descriptor == api_pb2.FILE_DESCRIPTOR_STDERR:
+            sr_fd = sr_pb2.TASK_EXEC_STDIO_FILE_DESCRIPTOR_STDERR
+        elif file_descriptor == api_pb2.FILE_DESCRIPTOR_INFO or file_descriptor == api_pb2.FILE_DESCRIPTOR_UNSPECIFIED:
+            raise ValueError(f"Unsupported file descriptor: {file_descriptor}")
+        else:
+            raise ValueError(f"Invalid file descriptor: {file_descriptor}")
+
+        async for item in self._stream_stdio(task_id, exec_id, sr_fd, deadline):
+            yield item
+
+    async def exec_stdin_write(
+        self, task_id: str, exec_id: str, offset: int, data: bytes, eof: bool
+    ) -> sr_pb2.TaskExecStdinWriteResponse:
+        """Write to the stdin stream of an exec'd command, properly retrying on transient errors.
+
+        Args:
+            task_id: The task ID of the task running the exec'd command.
+            exec_id: The execution ID of the command to write to.
+            offset: The offset to start writing to.
+            data: The data to write to the stdin stream.
+            eof: Whether to close the stdin stream after writing the data.
+        Raises:
+            Other errors: If retries are exhausted on transient errors or if there's an error
+              from the RPC itself.
+        """
+        request = sr_pb2.TaskExecStdinWriteRequest(task_id=task_id, exec_id=exec_id, offset=offset, data=data, eof=eof)
+        return await call_with_retries_on_transient_errors(
+            lambda: self._call_with_auth_retry(self._stub.TaskExecStdinWrite, request)
+        )
+
+    async def exec_poll(
+        self, task_id: str, exec_id: str, deadline: Optional[float] = None
+    ) -> sr_pb2.TaskExecPollResponse:
+        """Poll for the exit status of an exec'd command, properly retrying on transient errors.
+
+        Args:
+            task_id: The task ID of the task running the exec'd command.
+            exec_id: The execution ID of the command to poll on.
+        Returns:
+            sr_pb2.TaskExecPollResponse: The exit status of the command if it has completed.
+
+        Raises:
+            ExecTimeoutError: If the deadline is exceeded.
+            Other errors: If retries are exhausted on transient errors or if there's an error
+              from the RPC itself.
+        """
+        request = sr_pb2.TaskExecPollRequest(task_id=task_id, exec_id=exec_id)
+        # The timeout here is really a backstop in the event of a hang contacting
+        # the command router. Poll should usually be instantaneous.
+        timeout = deadline - time.monotonic() if deadline is not None else None
+        if timeout is not None and timeout <= 0:
+            raise ExecTimeoutError(f"Deadline exceeded while polling for exec {exec_id}")
+        try:
+            return await asyncio.wait_for(
+                call_with_retries_on_transient_errors(
+                    lambda: self._call_with_auth_retry(self._stub.TaskExecPoll, request)
+                ),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            raise ExecTimeoutError(f"Deadline exceeded while polling for exec {exec_id}")
+
+    async def exec_wait(
+        self,
+        task_id: str,
+        exec_id: str,
+        deadline: Optional[float] = None,
+    ) -> sr_pb2.TaskExecWaitResponse:
+        """Wait for an exec'd command to exit and return the exit code, properly retrying on transient errors.
+
+        Args:
+            task_id: The task ID of the task running the exec'd command.
+            exec_id: The execution ID of the command to wait on.
+        Returns:
+            Optional[sr_pb2.TaskExecWaitResponse]: The exit code of the command.
+        Raises:
+            ExecTimeoutError: If the deadline is exceeded.
+            Other errors: If there's an error from the RPC itself.
+        """
+        request = sr_pb2.TaskExecWaitRequest(task_id=task_id, exec_id=exec_id)
+        timeout = deadline - time.monotonic() if deadline is not None else None
+        if timeout is not None and timeout <= 0:
+            raise ExecTimeoutError(f"Deadline exceeded while waiting for exec {exec_id}")
+        try:
+            return await asyncio.wait_for(
+                call_with_retries_on_transient_errors(
+                    # We set a 60s timeout here to avoid waiting forever if there's an unanticipated hang
+                    # due to a networking issue. call_with_retries_on_transient_errors will retry if the
+                    # timeout is exceeded, so we'll retry every 60s until the command exits.
+                    #
+                    # Safety:
+                    # * If just the task shuts down, the task command router will return a NOT_FOUND error,
+                    #   and we'll stop retrying.
+                    # * If the task shut down AND the worker shut down, this could
+                    #   infinitely retry. For callers without an exec deadline, this
+                    #   could hang indefinitely.
+                    lambda: self._call_with_auth_retry(self._stub.TaskExecWait, request, timeout=60),
+                    base_delay_secs=1,  # Retry after 1s since total time is expected to be long.
+                    delay_factor=1,  # Fixed delay.
+                    max_retries=None,  # Retry forever.
+                ),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            raise ExecTimeoutError(f"Deadline exceeded while waiting for exec {exec_id}")
+
+    async def _refresh_jwt(self) -> None:
+        """Refresh JWT from the server and update internal state.
+
+        Concurrency-safe: only one refresh runs at a time.
+        """
+        async with self._jwt_refresh_lock:
+            if self._closed:
+                return
+
+            # If the current JWT expiration is already far enough in the future, don't refresh.
+            if self._jwt_exp is not None and self._jwt_exp - time.time() > 30:
+                # This can happen if multiple concurrent requests to the task command router
+                # get UNAUTHENTICATED errors and all refresh at the same time - one of them
+                # will win and the others will not refresh.
+                logger.debug(
+                    f"Skipping JWT refresh for exec with task ID {self._task_id} "
+                    "because its expiration is already far enough in the future"
+                )
+                return
+
+            resp = await fetch_command_router_access(self._server_client, self._task_id)
+            # Ensure the server URL remains stable for the lifetime of this client.
+            assert resp.url == self._server_url, "Task router URL changed during session"
+            self._jwt = resp.jwt
+            self._jwt_exp = _parse_jwt_expiration(resp.jwt)
+            # Wake up the background loop to recompute its next sleep.
+            self._jwt_refresh_event.set()
+
+    async def _call_with_auth_retry(self, func, *args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except GRPCError as exc:
+            if exc.status == Status.UNAUTHENTICATED:
+                await self._refresh_jwt()
+                # Retry with the original arguments preserved
+                return await func(*args, **kwargs)
+            raise
+
+    async def _jwt_refresh_loop(self) -> None:
+        """Background task that refreshes JWT 30 seconds before expiration.
+
+        Uses an event to wake early when a manual refresh happens or token changes.
+        """
+        while not self._closed:
+            try:
+                exp = self._jwt_exp
+                now = time.time()
+                if exp is None:
+                    # Unknown expiration: re-check periodically or until event wakes us.
+                    sleep_s = 60.0
+                else:
+                    refresh_at = exp - 30.0
+                    sleep_s = max(refresh_at - now, 0.0)
+
+                self._jwt_refresh_event.clear()
+                if sleep_s > 0:
+                    try:
+                        logger.debug(f"Waiting for JWT refresh for {sleep_s}s for exec with task ID {self._task_id}")
+                        # Wait until it's time to refresh, unless woken early.
+                        await asyncio.wait_for(self._jwt_refresh_event.wait(), timeout=sleep_s)
+                        logger.debug(f"Stopped waiting for JWT refresh for exec with task ID {self._task_id}")
+                        # Event fired (e.g., token changed) -> recompute timings.
+                        continue
+                    except asyncio.TimeoutError:
+                        logger.debug(f"Done waiting for JWT refresh for exec with task ID {self._task_id}")
+                        pass
+
+                # Time to refresh.
+                logger.debug(f"Refreshing JWT for exec with task ID {self._task_id}")
+                await self._refresh_jwt()
+            except asyncio.CancelledError:
+                logger.debug(f"Cancelled JWT refresh loop for exec with task ID {self._task_id}")
+                break
+            except Exception as e:
+                # Exceptions here can stem from non-transient errors against the server sending
+                # the TaskGetCommandRouterAccess RPC, for instance, if the task has finished.
+                logger.warning(f"Background JWT refresh failed for exec with task ID {self._task_id}: {e}")
+                break
+
+    async def _stream_stdio(
+        self,
+        task_id: str,
+        exec_id: str,
+        # Quotes around the type required for protobuf 3.19.
+        file_descriptor: "sr_pb2.TaskExecStdioFileDescriptor.ValueType",
+        deadline: Optional[float] = None,
+    ) -> AsyncIterator[sr_pb2.TaskExecStdioReadResponse]:
+        """Stream stdio from the task, properly updating the offset and retrying on transient errors.
+        Raises ExecTimeoutError if the deadline is exceeded.
+        """
+        offset = 0
+        delay_secs = self.stream_stdio_retry_delay_secs
+        delay_factor = self.stream_stdio_retry_delay_factor
+        num_retries_remaining = self.stream_stdio_max_retries
+        num_auth_retries = 0
+
+        async def sleep_and_update_delay_and_num_retries_remaining(e: Exception):
+            nonlocal delay_secs, num_retries_remaining
+            logger.debug(f"Retrying stdio read with delay {delay_secs}s due to error: {e}")
+            if deadline is not None and deadline - time.monotonic() <= delay_secs:
+                raise ExecTimeoutError(f"Deadline exceeded while streaming stdio for exec {exec_id}")
+
+            await asyncio.sleep(delay_secs)
+            delay_secs *= delay_factor
+            num_retries_remaining -= 1
+
+        while True:
+            timeout = max(0, deadline - time.monotonic()) if deadline is not None else None
+            try:
+                stream = self._stub.TaskExecStdioRead.open(timeout=timeout)
+                async with stream as s:
+                    req = sr_pb2.TaskExecStdioReadRequest(
+                        task_id=task_id,
+                        exec_id=exec_id,
+                        offset=offset,
+                        file_descriptor=file_descriptor,
+                    )
+
+                    # Scope auth retry strictly to the initial send (where headers/auth are sent).
+                    try:
+                        await s.send_message(req, end=True)
+                    except GRPCError as exc:
+                        if exc.status == Status.UNAUTHENTICATED and num_auth_retries < 1:
+                            await self._refresh_jwt()
+                            num_auth_retries += 1
+                            continue
+                        raise
+
+                    # We successfully authenticated, reset the auth retry count.
+                    num_auth_retries = 0
+
+                    async for item in s:
+                        # Reset retry backoff after any successful chunk.
+                        delay_secs = self.stream_stdio_retry_delay_secs
+                        offset += len(item.data)
+                        yield item
+
+                # We successfully streamed all output.
+                return
+            except GRPCError as e:
+                if num_retries_remaining > 0 and e.status in RETRYABLE_GRPC_STATUS_CODES:
+                    await sleep_and_update_delay_and_num_retries_remaining(e)
+                else:
+                    raise e
+            except AttributeError as e:
+                # StreamTerminatedError are not properly raised in grpclib<=0.4.7
+                # fixed in https://github.com/vmagamedov/grpclib/issues/185
+                # TODO: update to newer version (>=0.4.8) once stable
+                if num_retries_remaining > 0 and "_write_appdata" in str(e):
+                    await sleep_and_update_delay_and_num_retries_remaining(e)
+                else:
+                    raise e
+            except StreamTerminatedError as e:
+                if num_retries_remaining > 0:
+                    await sleep_and_update_delay_and_num_retries_remaining(e)
+                else:
+                    raise e
+            except asyncio.TimeoutError as e:
+                if num_retries_remaining > 0:
+                    await sleep_and_update_delay_and_num_retries_remaining(e)
+                else:
+                    raise ConnectionError(str(e))
+            except OSError as e:
+                if num_retries_remaining > 0:
+                    await sleep_and_update_delay_and_num_retries_remaining(e)
+                else:
+                    raise ConnectionError(str(e))