modal 1.0.3.dev10__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/config.py

@@ -66,11 +66,15 @@ Other possible configuration options are:
 * `traceback` (in the .toml file) / `MODAL_TRACEBACK` (as an env var).
   Defaults to False. Enables printing full tracebacks on unexpected CLI
   errors, which can be useful for debugging client issues.
-* `log_pattern` (in the .toml file) / MODAL_LOG_PATTERN` (as an env var).
-  Defaults to "[modal-client] %(asctime)s %(message)s"
+* `log_pattern` (in the .toml file) / `MODAL_LOG_PATTERN` (as an env var).
+  Defaults to `"[modal-client] %(asctime)s %(message)s"`
   The log formatting pattern that will be used by the modal client itself.
   See https://docs.python.org/3/library/logging.html#logrecord-attributes for available
   log attributes.
+* `dev_suffix` (in the .toml file) / `MODAL_DEV_SUFFIX` (as an env var).
+  Overrides the default `-dev` suffix added to URLs generated for web endpoints
+  when the App is ephemeral (i.e., created via `modal serve`). Must be a short
+  alphanumeric string.
 
 Meta-configuration
 ------------------
@@ -85,6 +89,7 @@ Some "meta-options" are set using environment variables only:
 
 import logging
 import os
+import re
 import typing
 import warnings
 from typing import Any, Callable, Optional
@@ -94,7 +99,7 @@ from google.protobuf.empty_pb2 import Empty
 from modal_proto import api_pb2
 
 from ._utils.logger import configure_logger
-from .exception import InvalidError
+from .exception import InvalidError, NotFoundError
 
 DEFAULT_SERVER_URL = "https://api.modal.com"
 
@@ -142,7 +147,7 @@ async def _lookup_workspace(server_url: str, token_id: str, token_secret: str) -
 
     credentials = (token_id, token_secret)
     async with _Client(server_url, api_pb2.CLIENT_TYPE_CLIENT, credentials) as client:
-        return await client.stub.WorkspaceNameLookup(Empty(), timeout=3)
+        return await client.stub.WorkspaceNameLookup(Empty(), retry=None, timeout=3)
 
 
 def config_profiles():
@@ -158,15 +163,15 @@ def _config_active_profile() -> str:
         return "default"
 
 
-def config_set_active_profile(env: str) -> None:
+def config_set_active_profile(profile: str) -> None:
     """Set the user's active modal profile by writing it to the `.modal.toml` file."""
-    if env not in _user_config:
-        raise KeyError(env)
+    if profile not in _user_config:
+        raise NotFoundError(f"No profile named '{profile}' found in {user_config_path}")
 
-    for key, values in _user_config.items():
-        values.pop("active", None)
+    for profile_data in _user_config.values():
+        profile_data.pop("active", None)
 
-    _user_config[env]["active"] = True
+    _user_config[profile]["active"] = True  # type: ignore
     _write_user_config(_user_config)
 
 
@@ -206,6 +211,12 @@ def _check_value(options: list[str]) -> Callable[[str], str]:
     return checker
 
 
+def _enforce_suffix_rules(x: str) -> str:
+    if x and not re.match(r"^[a-zA-Z0-9]{1,8}$", x):
+        raise ValueError("Suffix must be an alphanumeric string of no more than 8 characters.")
+    return x
+
+
 class _Setting(typing.NamedTuple):
     default: typing.Any = None
     transform: typing.Callable[[str], typing.Any] = lambda x: x  # noqa: E731
@@ -236,10 +247,17 @@ _SETTINGS = {
     "traceback": _Setting(False, transform=_to_boolean),
     "image_builder_version": _Setting(),
     "strict_parameters": _Setting(False, transform=_to_boolean),  # For internal/experimental use
+    # Allow insecure TLS for the task command router when running locally (testing/dev only)
+    "task_command_router_insecure": _Setting(False, transform=_to_boolean),
     "snapshot_debug": _Setting(False, transform=_to_boolean),
     "cuda_checkpoint_path": _Setting("/__modal/.bin/cuda-checkpoint"),  # Used for snapshotting GPU memory.
-    "function_schemas": _Setting(False, transform=_to_boolean),
     "build_validation": _Setting("error", transform=_check_value(["error", "warn", "ignore"])),
+    # Payload format for function inputs/outputs: 'pickle' (default) or 'cbor'
+    "payload_format": _Setting(
+        "pickle",
+        transform=lambda s: _check_value(["pickle", "cbor"])(s.lower()),
+    ),
+    "dev_suffix": _Setting("", transform=_enforce_suffix_rules),
 }
 
 

modal/container_process.py

@@ -1,26 +1,29 @@
 # Copyright Modal Labs 2024
 import asyncio
 import platform
+import time
 from typing import Generic, Optional, TypeVar
 
 from modal_proto import api_pb2
 
 from ._utils.async_utils import TaskContext, synchronize_api
-from ._utils.grpc_utils import retry_transient_errors
 from ._utils.shell_utils import stream_from_stdin, write_to_fd
+from ._utils.task_command_router_client import TaskCommandRouterClient
 from .client import _Client
-from .exception import InteractiveTimeoutError, InvalidError
+from .config import logger
+from .exception import ExecTimeoutError, InteractiveTimeoutError, InvalidError
 from .io_streams import _StreamReader, _StreamWriter
 from .stream_type import StreamType
 
 T = TypeVar("T", str, bytes)
 
 
-class _ContainerProcess(Generic[T]):
+class _ContainerProcessThroughServer(Generic[T]):
     _process_id: Optional[str] = None
     _stdout: _StreamReader[T]
     _stderr: _StreamReader[T]
     _stdin: _StreamWriter
+    _exec_deadline: Optional[float] = None
     _text: bool
     _by_line: bool
     _returncode: Optional[int] = None
@@ -28,14 +31,17 @@ class _ContainerProcess(Generic[T]):
     def __init__(
         self,
         process_id: str,
+        task_id: str,
         client: _Client,
         stdout: StreamType = StreamType.PIPE,
         stderr: StreamType = StreamType.PIPE,
+        exec_deadline: Optional[float] = None,
         text: bool = True,
         by_line: bool = False,
     ) -> None:
         self._process_id = process_id
         self._client = client
+        self._exec_deadline = exec_deadline
         self._text = text
         self._by_line = by_line
         self._stdout = _StreamReader[T](
@@ -46,6 +52,8 @@ class _ContainerProcess(Generic[T]):
             stream_type=stdout,
             text=text,
             by_line=by_line,
+            deadline=exec_deadline,
+            task_id=task_id,
         )
         self._stderr = _StreamReader[T](
             api_pb2.FILE_DESCRIPTOR_STDERR,
@@ -55,6 +63,8 @@ class _ContainerProcess(Generic[T]):
             stream_type=stderr,
             text=text,
             by_line=by_line,
+            deadline=exec_deadline,
+            task_id=task_id,
         )
         self._stdin = _StreamWriter(process_id, "container_process", self._client)
 
@@ -90,11 +100,17 @@ class _ContainerProcess(Generic[T]):
 
         Returns `None` if the process is still running, else returns the exit code.
         """
+        assert self._process_id
         if self._returncode is not None:
             return self._returncode
+        if self._exec_deadline and time.monotonic() >= self._exec_deadline:
+            # TODO(matt): In the future, it would be nice to raise a ContainerExecTimeoutError to make it
+            # clear to the user that their sandbox terminated due to a timeout
+            self._returncode = -1
+            return self._returncode
 
         req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=0)
-        resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(self._client.stub.ContainerExecWait, req)
+        resp = await self._client.stub.ContainerExecWait(req)
 
         if resp.completed:
             self._returncode = resp.exit_code
@@ -102,40 +118,56 @@ class _ContainerProcess(Generic[T]):
 
         return None
 
+    async def _wait_for_completion(self) -> int:
+        assert self._process_id
+        while True:
+            req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=10)
+            resp = await self._client.stub.ContainerExecWait(req)
+            if resp.completed:
+                return resp.exit_code
+
     async def wait(self) -> int:
         """Wait for the container process to finish running. Returns the exit code."""
-
         if self._returncode is not None:
             return self._returncode
 
-        while True:
-            req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=10)
-            resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(
-                self._client.stub.ContainerExecWait, req
-            )
-            if resp.completed:
-                self._returncode = resp.exit_code
-                return self._returncode
+        try:
+            timeout = None
+            if self._exec_deadline:
+                timeout = self._exec_deadline - time.monotonic()
+                if timeout <= 0:
+                    raise TimeoutError()
+            self._returncode = await asyncio.wait_for(self._wait_for_completion(), timeout=timeout)
+        except (asyncio.TimeoutError, TimeoutError):
+            self._returncode = -1
+        logger.debug(f"ContainerProcess {self._process_id} wait completed with returncode {self._returncode}")
+        return self._returncode
 
     async def attach(self):
+        """mdmd:hidden"""
         if platform.system() == "Windows":
             print("interactive exec is not currently supported on Windows.")
             return
 
-        from rich.console import Console
+        from ._output import make_console
 
-        console = Console()
+        console = make_console()
 
         connecting_status = console.status("Connecting...")
         connecting_status.start()
         on_connect = asyncio.Event()
 
         async def _write_to_fd_loop(stream: _StreamReader):
+            # This is required to make modal shell to an existing task work,
+            # since that uses ContainerExec RPCs directly, but this is hacky.
+            #
+            # TODO(saltzm): Once we use the new exec path for that use case, this code can all be removed.
+            from .io_streams import _StreamReaderThroughServer
+
+            assert isinstance(stream._impl, _StreamReaderThroughServer)
+            stream_impl = stream._impl
             # Don't skip empty messages so we can detect when the process has booted.
-            async for chunk in stream._get_logs(skip_empty_messages=False):
-                if chunk is None:
-                    break
-
+            async for chunk in stream_impl._get_logs(skip_empty_messages=False):
                 if not on_connect.is_set():
                     connecting_status.stop()
                     on_connect.set()
@@ -169,4 +201,272 @@ class _ContainerProcess(Generic[T]):
                 raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
 
 
+async def _iter_stream_as_bytes(stream: _StreamReader[T]):
+    """Yield raw bytes from a StreamReader regardless of text mode/backend."""
+    async for part in stream:
+        if isinstance(part, str):
+            yield part.encode("utf-8")
+        else:
+            yield part
+
+
+class _ContainerProcessThroughCommandRouter(Generic[T]):
+    """
+    Container process implementation that works via direct communication with
+    the Modal worker where the container is running.
+    """
+
+    def __init__(
+        self,
+        process_id: str,
+        client: _Client,
+        command_router_client: TaskCommandRouterClient,
+        task_id: str,
+        *,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        exec_deadline: Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        self._client = client
+        self._command_router_client = command_router_client
+        self._process_id = process_id
+        self._exec_deadline = exec_deadline
+        self._text = text
+        self._by_line = by_line
+        self._task_id = task_id
+        self._stdout = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDOUT,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stdout,
+            text=text,
+            by_line=by_line,
+            deadline=exec_deadline,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._stderr = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDERR,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stderr,
+            text=text,
+            by_line=by_line,
+            deadline=exec_deadline,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._stdin = _StreamWriter(
+            process_id,
+            "container_process",
+            self._client,
+            command_router_client=self._command_router_client,
+            task_id=self._task_id,
+        )
+        self._returncode = None
+
+    def __repr__(self) -> str:
+        return f"ContainerProcess(process_id={self._process_id!r})"
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        return self._stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        return self._stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        return self._stdin
+
+    @property
+    def returncode(self) -> int:
+        if self._returncode is None:
+            raise InvalidError(
+                "You must call wait() before accessing the returncode. "
+                "To poll for the status of a running process, use poll() instead."
+            )
+        return self._returncode
+
+    async def poll(self) -> Optional[int]:
+        if self._returncode is not None:
+            return self._returncode
+        try:
+            resp = await self._command_router_client.exec_poll(self._task_id, self._process_id, self._exec_deadline)
+            which = resp.WhichOneof("exit_status")
+            if which is None:
+                return None
+
+            if which == "code":
+                self._returncode = int(resp.code)
+                return self._returncode
+            elif which == "signal":
+                self._returncode = 128 + int(resp.signal)
+                return self._returncode
+            else:
+                logger.debug(f"ContainerProcess {self._process_id} exited with unexpected status: {which}")
+                raise InvalidError("Unexpected exit status")
+        except ExecTimeoutError:
+            logger.debug(f"ContainerProcess poll for {self._process_id} did not complete within deadline")
+            # TODO(saltzm): This is a weird API, but customers currently may rely on it. This
+            # should probably raise an ExecTimeoutError instead.
+            self._returncode = -1
+            return self._returncode
+        except Exception as e:
+            # Re-raise non-transient errors or errors resulting from exceeding retries on transient errors.
+            logger.warning(f"ContainerProcess poll for {self._process_id} failed: {e}")
+            raise
+
+    async def wait(self) -> int:
+        if self._returncode is not None:
+            return self._returncode
+
+        try:
+            resp = await self._command_router_client.exec_wait(self._task_id, self._process_id, self._exec_deadline)
+            which = resp.WhichOneof("exit_status")
+            if which == "code":
+                self._returncode = int(resp.code)
+            elif which == "signal":
+                self._returncode = 128 + int(resp.signal)
+            else:
+                logger.debug(f"ContainerProcess {self._process_id} exited with unexpected status: {which}")
+                self._returncode = -1
+                raise InvalidError("Unexpected exit status")
+        except ExecTimeoutError:
+            logger.debug(f"ContainerProcess {self._process_id} did not complete within deadline")
+            # TODO(saltzm): This is a weird API, but customers currently may rely on it. This
+            # should be a ExecTimeoutError.
+            self._returncode = -1
+
+        return self._returncode
+
+    async def attach(self):
+        if platform.system() == "Windows":
+            print("interactive exec is not currently supported on Windows.")
+            return
+
+        from ._output import make_console
+
+        console = make_console()
+
+        connecting_status = console.status("Connecting...")
+        connecting_status.start()
+        on_connect = asyncio.Event()
+
+        async def _write_to_fd_loop(stream: _StreamReader[T]):
+            async for chunk in _iter_stream_as_bytes(stream):
+                if chunk is None:
+                    break
+
+                if not on_connect.is_set():
+                    connecting_status.stop()
+                    on_connect.set()
+
+                await write_to_fd(stream.file_descriptor, chunk)
+
+        async def _handle_input(data: bytes, message_index: int):
+            self.stdin.write(data)
+            await self.stdin.drain()
+
+        async with TaskContext() as tc:
+            stdout_task = tc.create_task(_write_to_fd_loop(self.stdout))
+            stderr_task = tc.create_task(_write_to_fd_loop(self.stderr))
+
+            try:
+                # Time out if we can't connect fast enough.
+                await asyncio.wait_for(on_connect.wait(), timeout=60)
+
+                async with stream_from_stdin(_handle_input, use_raw_terminal=True):
+                    await stdout_task
+                    await stderr_task
+
+            except (asyncio.TimeoutError, TimeoutError):
+                connecting_status.stop()
+                stdout_task.cancel()
+                stderr_task.cancel()
+                raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
+
+
+class _ContainerProcess(Generic[T]):
+    """Represents a running process in a container."""
+
+    def __init__(
+        self,
+        process_id: str,
+        task_id: str,
+        client: _Client,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        exec_deadline: Optional[float] = None,
+        text: bool = True,
+        by_line: bool = False,
+        command_router_client: Optional[TaskCommandRouterClient] = None,
+    ) -> None:
+        if command_router_client is None:
+            self._impl = _ContainerProcessThroughServer(
+                process_id,
+                task_id,
+                client,
+                stdout=stdout,
+                stderr=stderr,
+                exec_deadline=exec_deadline,
+                text=text,
+                by_line=by_line,
+            )
+        else:
+            self._impl = _ContainerProcessThroughCommandRouter(
+                process_id,
+                client,
+                command_router_client,
+                task_id,
+                stdout=stdout,
+                stderr=stderr,
+                exec_deadline=exec_deadline,
+                text=text,
+                by_line=by_line,
+            )
+
+    def __repr__(self) -> str:
+        return self._impl.__repr__()
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        return self._impl.stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        return self._impl.stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        return self._impl.stdin
+
+    @property
+    def returncode(self) -> int:
+        return self._impl.returncode
+
+    async def poll(self) -> Optional[int]:
+        """Check if the container process has finished running.
+
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        return await self._impl.poll()
+
+    async def wait(self) -> int:
+        """Wait for the container process to finish running. Returns the exit code."""
+        return await self._impl.wait()
+
+    async def attach(self):
+        """mdmd:hidden"""
+        await self._impl.attach()
+
+
 ContainerProcess = synchronize_api(_ContainerProcess)