modal 1.1.5.dev83__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/_utils/task_command_router_client.py

@@ -0,0 +1,524 @@
+# Copyright Modal Labs 2025
+import asyncio
+import base64
+import json
+import ssl
+import time
+import urllib.parse
+import weakref
+from contextlib import suppress
+from typing import AsyncGenerator, 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 ConflictError, 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_client import grpc_error_converter
+from .async_utils import aclosing
+from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES, connect_channel
+
+
+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 server_client.stub.TaskGetCommandRouterAccess(
+        api_pb2.TaskGetCommandRouterAccessRequest(task_id=task_id),
+    )
+
+
+def _finalize_channel(loop, channel):
+    if not loop.is_closed():
+        # only run if loop has not shut down
+        # call_soon_threadsafe could throw if the loop is torn down after calling
+        # is_closed. This is safe to ignore.
+        with suppress(Exception):
+            loop.call_soon_threadsafe(channel.close)
+
+
+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 ConflictError:
+            logger.debug(f"Command router access is not enabled for task {task_id}")
+            return None
+
+        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)
+        loop = asyncio.get_running_loop()
+        jwt_refresh_lock = asyncio.Lock()
+
+        return cls(server_client, task_id, resp.url, resp.jwt, channel, loop, jwt_refresh_lock)
+
+    def __init__(
+        self,
+        server_client,
+        task_id: str,
+        server_url: str,
+        jwt: str,
+        channel: grpclib.client.Channel,
+        loop: asyncio.AbstractEventLoop,
+        jwt_refresh_lock: asyncio.Lock,
+        *,
+        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."""
+        # Record the loop this instance is bound to so __del__ can safely schedule cleanup
+        # even if finalization happens from a different thread (e.g. via synchronicity).
+        self._loop = loop
+
+        # 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)
+        # This is passed in as an argument to ensure it's created from within the correct event loop.
+        self._jwt_refresh_lock = jwt_refresh_lock
+
+        self._closed = False
+
+        self._channel_finalizer = weakref.finalize(
+            self,
+            _finalize_channel,
+            loop,
+            channel,
+        )
+
+        self._stub = TaskCommandRouterStub(self._channel)
+
+    def _get_metadata(self):
+        return {"authorization": f"Bearer {self._jwt}"}
+
+    async def close(self) -> None:
+        """Close the client."""
+        if self._closed:
+            return
+
+        self._closed = True
+        self._channel.close()
+        if self._channel_finalizer.alive:
+            # skip the finalizer if we've closed the channel anyway
+            self._channel_finalizer.detach()
+
+    async def exec_start(self, request: sr_pb2.TaskExecStartRequest) -> sr_pb2.TaskExecStartResponse:
+        """Start an exec'd command, properly retrying on transient errors."""
+        with grpc_error_converter():
+            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,
+    ) -> AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]:
+        """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:
+            AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]: 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}")
+
+        with grpc_error_converter():
+            async with aclosing(self._stream_stdio(task_id, exec_id, sr_fd, deadline)) as stream:
+                async for item in stream:
+                    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)
+        with grpc_error_converter():
+            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}")
+
+        with grpc_error_converter():
+            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}")
+
+        with grpc_error_converter():
+            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."""
+        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
+
+            logger.debug(f"Refreshing JWT for exec with task ID {self._task_id}")
+            resp = await fetch_command_router_access(self._server_client, self._task_id)
+            logger.debug(f"Finished refreshing JWT for exec with task ID {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)
+
+    async def _call_with_auth_retry(self, func, *args, **kwargs):
+        try:
+            return await func(*args, **kwargs, metadata=self._get_metadata())
+        except GRPCError as exc:
+            if exc.status == Status.UNAUTHENTICATED:
+                await self._refresh_jwt()
+                # Retry with the original arguments preserved
+                return await func(*args, **kwargs, metadata=self._get_metadata())
+            raise
+
+    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,
+    ) -> AsyncGenerator[sr_pb2.TaskExecStdioReadResponse, None]:
+        """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
+        # Flag to prevent infinite auth retries in the event that the JWT
+        # refresh yields an invalid JWT somehow or that the JWT is otherwise invalid.
+        did_auth_retry = False
+
+        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, metadata=self._get_metadata())
+                async with stream as s:
+                    req = sr_pb2.TaskExecStdioReadRequest(
+                        task_id=task_id,
+                        exec_id=exec_id,
+                        offset=offset,
+                        file_descriptor=file_descriptor,
+                    )
+
+                    # Auth retry is scoped to a single refresh per streaming attempt. While auth metadata is
+                    # sent on request start, UNAUTHENTICATED may sometimes surface during iteration,
+                    # so we handle it at both send and receive boundaries.
+                    try:
+                        await s.send_message(req, end=True)
+                        async for item in s:
+                            # We successfully authenticated after a JWT refresh, reset the auth retry flag.
+                            if did_auth_retry:
+                                did_auth_retry = False
+                            # Reset retry backoff after any successful chunk.
+                            delay_secs = self.stream_stdio_retry_delay_secs
+                            offset += len(item.data)
+                            yield item
+                    except GRPCError as exc:
+                        if exc.status == Status.UNAUTHENTICATED and not did_auth_retry:
+                            await self._refresh_jwt()
+                            # Mark that we've retried authentication for this streaming attempt, to
+                            # prevent subsequent retries.
+                            did_auth_retry = True
+                            continue
+                        raise
+
+                # 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))
+
+    async def mount_image(self, request: sr_pb2.TaskMountDirectoryRequest):
+        with grpc_error_converter():
+            return await call_with_retries_on_transient_errors(
+                lambda: self._call_with_auth_retry(self._stub.TaskMountDirectory, request)
+            )
+
+    async def snapshot_directory(
+        self, request: sr_pb2.TaskSnapshotDirectoryRequest
+    ) -> sr_pb2.TaskSnapshotDirectoryResponse:
+        with grpc_error_converter():
+            return await call_with_retries_on_transient_errors(
+                lambda: self._call_with_auth_retry(self._stub.TaskSnapshotDirectory, request)
+            )