modal 0.62.16__py3-none-any.whl → 0.72.11__py3-none-any.whl

modal/container_process.py

@@ -0,0 +1,177 @@
+# Copyright Modal Labs 2024
+import asyncio
+import platform
+from typing import Generic, Optional, TypeVar
+
+from modal_proto import api_pb2
+
+from ._utils.async_utils import TaskContext, synchronize_api
+from ._utils.deprecation import deprecation_error
+from ._utils.grpc_utils import retry_transient_errors
+from ._utils.shell_utils import stream_from_stdin, write_to_fd
+from .client import _Client
+from .exception import InteractiveTimeoutError, InvalidError
+from .io_streams import _StreamReader, _StreamWriter
+from .stream_type import StreamType
+
+T = TypeVar("T", str, bytes)
+
+
+class _ContainerProcess(Generic[T]):
+    _process_id: Optional[str] = None
+    _stdout: _StreamReader[T]
+    _stderr: _StreamReader[T]
+    _stdin: _StreamWriter
+    _text: bool
+    _by_line: bool
+    _returncode: Optional[int] = None
+
+    def __init__(
+        self,
+        process_id: str,
+        client: _Client,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        self._process_id = process_id
+        self._client = client
+        self._text = text
+        self._by_line = by_line
+        self._stdout = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDOUT,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stdout,
+            text=text,
+            by_line=by_line,
+        )
+        self._stderr = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDERR,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stderr,
+            text=text,
+            by_line=by_line,
+        )
+        self._stdin = _StreamWriter(process_id, "container_process", self._client)
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        return self._stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        return self._stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        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]:
+        """Check if the container process has finished running.
+
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        if self._returncode is not None:
+            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)
+
+        if resp.completed:
+            self._returncode = resp.exit_code
+            return self._returncode
+
+        return None
+
+    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=50)
+            resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(
+                self._client.stub.ContainerExecWait, req
+            )
+            if resp.completed:
+                self._returncode = resp.exit_code
+                return self._returncode
+
+    async def attach(self, *, pty: Optional[bool] = None):
+        if platform.system() == "Windows":
+            print("interactive exec is not currently supported on Windows.")
+            return
+
+        if pty is not None:
+            deprecation_error(
+                (2024, 12, 9),
+                "The `pty` argument to `modal.container_process.attach(pty=...)` is deprecated, "
+                "as only PTY mode is supported. Please remove the argument.",
+            )
+
+        from rich.console import Console
+
+        console = Console()
+
+        connecting_status = console.status("Connecting...")
+        connecting_status.start()
+        on_connect = asyncio.Event()
+
+        async def _write_to_fd_loop(stream: _StreamReader):
+            # 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
+
+                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 to the server 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
+
+                # TODO: this doesn't work right now.
+                # if exit_status != 0:
+                #     raise ExecutionError(f"Process exited with status code {exit_status}")
+
+            except (asyncio.TimeoutError, TimeoutError):
+                connecting_status.stop()
+                stdout_task.cancel()
+                stderr_task.cancel()
+                raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
+
+
+ContainerProcess = synchronize_api(_ContainerProcess)

modal/container_process.pyi

@@ -0,0 +1,82 @@
+import modal.client
+import modal.io_streams
+import modal.stream_type
+import typing
+import typing_extensions
+
+T = typing.TypeVar("T")
+
+class _ContainerProcess(typing.Generic[T]):
+    _process_id: typing.Optional[str]
+    _stdout: modal.io_streams._StreamReader[T]
+    _stderr: modal.io_streams._StreamReader[T]
+    _stdin: modal.io_streams._StreamWriter
+    _text: bool
+    _by_line: bool
+    _returncode: typing.Optional[int]
+
+    def __init__(
+        self,
+        process_id: str,
+        client: modal.client._Client,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None: ...
+    @property
+    def stdout(self) -> modal.io_streams._StreamReader[T]: ...
+    @property
+    def stderr(self) -> modal.io_streams._StreamReader[T]: ...
+    @property
+    def stdin(self) -> modal.io_streams._StreamWriter: ...
+    @property
+    def returncode(self) -> int: ...
+    async def poll(self) -> typing.Optional[int]: ...
+    async def wait(self) -> int: ...
+    async def attach(self, *, pty: typing.Optional[bool] = None): ...
+
+class ContainerProcess(typing.Generic[T]):
+    _process_id: typing.Optional[str]
+    _stdout: modal.io_streams.StreamReader[T]
+    _stderr: modal.io_streams.StreamReader[T]
+    _stdin: modal.io_streams.StreamWriter
+    _text: bool
+    _by_line: bool
+    _returncode: typing.Optional[int]
+
+    def __init__(
+        self,
+        process_id: str,
+        client: modal.client.Client,
+        stdout: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        stderr: modal.stream_type.StreamType = modal.stream_type.StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None: ...
+    @property
+    def stdout(self) -> modal.io_streams.StreamReader[T]: ...
+    @property
+    def stderr(self) -> modal.io_streams.StreamReader[T]: ...
+    @property
+    def stdin(self) -> modal.io_streams.StreamWriter: ...
+    @property
+    def returncode(self) -> int: ...
+
+    class __poll_spec(typing_extensions.Protocol):
+        def __call__(self) -> typing.Optional[int]: ...
+        async def aio(self) -> typing.Optional[int]: ...
+
+    poll: __poll_spec
+
+    class __wait_spec(typing_extensions.Protocol):
+        def __call__(self) -> int: ...
+        async def aio(self) -> int: ...
+
+    wait: __wait_spec
+
+    class __attach_spec(typing_extensions.Protocol):
+        def __call__(self, *, pty: typing.Optional[bool] = None): ...
+        async def aio(self, *, pty: typing.Optional[bool] = None): ...
+
+    attach: __attach_spec

modal/dict.py

@@ -1,6 +1,8 @@
 # Copyright Modal Labs 2022
-from typing import Any, AsyncIterator, Optional, Type
+from collections.abc import AsyncIterator
+from typing import Any, Optional
 
+from grpclib import GRPCError
 from synchronicity.async_wrap import asynccontextmanager
 
 from modal_proto import api_pb2
@@ -8,11 +10,13 @@ from modal_proto import api_pb2
 from ._resolver import Resolver
 from ._serialization import deserialize, serialize
 from ._utils.async_utils import TaskContext, synchronize_api
+from ._utils.deprecation import renamed_parameter
 from ._utils.grpc_utils import retry_transient_errors
+from ._utils.name_utils import check_object_name
 from .client import _Client
 from .config import logger
-from .exception import deprecation_warning
-from .object import EPHEMERAL_OBJECT_HEARTBEAT_SLEEP, _get_environment_name, _Object, live_method
+from .exception import RequestSizeError
+from .object import EPHEMERAL_OBJECT_HEARTBEAT_SLEEP, _get_environment_name, _Object, live_method, live_method_gen
 
 
 def _serialize_dict(data):
@@ -22,54 +26,38 @@ def _serialize_dict(data):
 class _Dict(_Object, type_prefix="di"):
     """Distributed dictionary for storage in Modal apps.
 
-    Keys and values can be essentially any object, so long as they can be
-    serialized by `cloudpickle`, including Modal objects.
+    Keys and values can be essentially any object, so long as they can be serialized by
+    `cloudpickle`, which includes other Modal objects.
 
     **Lifetime of a Dict and its items**
 
     An individual dict entry will expire 30 days after it was last added to its Dict object.
-    Because of this, `Dict`s are best not used for
-    long-term storage. All data is deleted when the app is stopped.
+    Additionally, data are stored in memory on the Modal server and could be lost due to
+    unexpected server restarts. Because of this, `Dict` is best suited for storing short-term
+    state and is not recommended for durable storage.
 
     **Usage**
 
     ```python
-    from modal import Dict, Stub
+    from modal import Dict
 
-    stub = Stub()
     my_dict = Dict.from_name("my-persisted_dict", create_if_missing=True)
 
-    @stub.local_entrypoint()
-    def main():
-        my_dict["some key"] = "some value"
-        my_dict[123] = 456
+    my_dict["some key"] = "some value"
+    my_dict[123] = 456
 
-        assert my_dict["some key"] == "some value"
-        assert my_dict[123] == 456
+    assert my_dict["some key"] == "some value"
+    assert my_dict[123] == 456
     ```
 
+    The `Dict` class offers a few methods for operations that are usually accomplished
+    in Python with operators, such as `Dict.put` and `Dict.contains`. The advantage of
+    these methods is that they can be safely called in an asynchronous context, whereas
+    their operator-based analogues will block the event loop.
+
     For more examples, see the [guide](/docs/guide/dicts-and-queues#modal-dicts).
     """
 
-    @staticmethod
-    def new(data: Optional[dict] = None) -> "_Dict":
-        """`Dict.new` is deprecated.
-
-        Please use `Dict.from_name` (for persisted) or `Dict.ephemeral` (for ephemeral) dicts.
-        """
-        deprecation_warning((2024, 3, 19), Dict.new.__doc__)
-
-        async def _load(self: _Dict, resolver: Resolver, existing_object_id: Optional[str]):
-            serialized = _serialize_dict(data if data is not None else {})
-            req = api_pb2.DictCreateRequest(
-                app_id=resolver.app_id, data=serialized, existing_dict_id=existing_object_id
-            )
-            response = await resolver.client.stub.DictCreate(req)
-            logger.debug(f"Created dict with id {response.dict_id}")
-            self._hydrate(response.dict_id, resolver.client, None)
-
-        return _Dict._from_loader(_load, "Dict()")
-
     def __init__(self, data={}):
         """mdmd:hidden"""
         raise RuntimeError(
@@ -79,7 +67,7 @@ class _Dict(_Object, type_prefix="di"):
     @classmethod
     @asynccontextmanager
     async def ephemeral(
-        cls: Type["_Dict"],
+        cls: type["_Dict"],
         data: Optional[dict] = None,
         client: Optional[_Client] = None,
         environment_name: Optional[str] = None,
@@ -89,9 +77,13 @@ class _Dict(_Object, type_prefix="di"):
 
         Usage:
         ```python
+        from modal import Dict
+
         with Dict.ephemeral() as d:
             d["foo"] = "bar"
+        ```
 
+        ```python notest
         async with Dict.ephemeral() as d:
             await d.put.aio("foo", "bar")
         ```
@@ -104,37 +96,38 @@ class _Dict(_Object, type_prefix="di"):
             environment_name=_get_environment_name(environment_name),
             data=serialized,
         )
-        response = await client.stub.DictGetOrCreate(request)
+        response = await retry_transient_errors(client.stub.DictGetOrCreate, request, total_timeout=10.0)
         async with TaskContext() as tc:
             request = api_pb2.DictHeartbeatRequest(dict_id=response.dict_id)
             tc.infinite_loop(lambda: client.stub.DictHeartbeat(request), sleep=_heartbeat_sleep)
             yield cls._new_hydrated(response.dict_id, client, None, is_another_app=True)
 
     @staticmethod
+    @renamed_parameter((2024, 12, 18), "label", "name")
     def from_name(
-        label: str,
+        name: str,
         data: Optional[dict] = None,
         namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
         environment_name: Optional[str] = None,
         create_if_missing: bool = False,
     ) -> "_Dict":
-        """Create a reference to a persisted Dict
-
-        **Examples**
+        """Reference a named Dict, creating if necessary.
 
-        ```python notest
-        # In one app:
-        stub.dict = Dict.persisted("my-dict")
+        In contrast to `modal.Dict.lookup`, this is a lazy method
+        that defers hydrating the local object with metadata from
+        Modal servers until the first time it is actually used.
 
-        # Later, in another app or Python file:
-        stub.dict = Dict.from_name("my-dict")
+        ```python
+        d = modal.Dict.from_name("my-dict", create_if_missing=True)
+        d[123] = 456
         ```
         """
+        check_object_name(name, "Dict")
 
         async def _load(self: _Dict, resolver: Resolver, existing_object_id: Optional[str]):
             serialized = _serialize_dict(data if data is not None else {})
             req = api_pb2.DictGetOrCreateRequest(
-                deployment_name=label,
+                deployment_name=name,
                 namespace=namespace,
                 environment_name=_get_environment_name(environment_name, resolver),
                 object_creation_type=(api_pb2.OBJECT_CREATION_TYPE_CREATE_IF_MISSING if create_if_missing else None),
@@ -144,26 +137,22 @@ class _Dict(_Object, type_prefix="di"):
             logger.debug(f"Created dict with id {response.dict_id}")
             self._hydrate(response.dict_id, resolver.client, None)
 
-        return _Dict._from_loader(_load, "Dict()", is_another_app=True)
-
-    @staticmethod
-    def persisted(
-        label: str, namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE, environment_name: Optional[str] = None
-    ) -> "_Dict":
-        """Deprecated! Use `Dict.from_name(name, create_if_missing=True)`."""
-        deprecation_warning((2024, 3, 1), _Dict.persisted.__doc__)
-        return _Dict.from_name(label, namespace, environment_name, create_if_missing=True)
+        return _Dict._from_loader(_load, "Dict()", is_another_app=True, hydrate_lazily=True)
 
     @staticmethod
+    @renamed_parameter((2024, 12, 18), "label", "name")
     async def lookup(
-        label: str,
+        name: str,
         data: Optional[dict] = None,
         namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
         client: Optional[_Client] = None,
         environment_name: Optional[str] = None,
         create_if_missing: bool = False,
     ) -> "_Dict":
-        """Lookup a dict with a given name and tag.
+        """Lookup a named Dict.
+
+        In contrast to `modal.Dict.from_name`, this is an eager method
+        that will hydrate the local object with metadata from Modal servers.
 
         ```python
         d = modal.Dict.lookup("my-dict")
@@ -171,7 +160,7 @@ class _Dict(_Object, type_prefix="di"):
         ```
         """
         obj = _Dict.from_name(
-            label,
+            name,
             data=data,
             namespace=namespace,
             environment_name=environment_name,
@@ -183,9 +172,21 @@ class _Dict(_Object, type_prefix="di"):
         await resolver.load(obj)
         return obj
 
+    @staticmethod
+    @renamed_parameter((2024, 12, 18), "label", "name")
+    async def delete(
+        name: str,
+        *,
+        client: Optional[_Client] = None,
+        environment_name: Optional[str] = None,
+    ):
+        obj = await _Dict.lookup(name, client=client, environment_name=environment_name)
+        req = api_pb2.DictDeleteRequest(dict_id=obj.object_id)
+        await retry_transient_errors(obj._client.stub.DictDelete, req)
+
     @live_method
     async def clear(self) -> None:
-        """Remove all items from the modal.Dict."""
+        """Remove all items from the Dict."""
         req = api_pb2.DictClearRequest(dict_id=self.object_id)
         await retry_transient_errors(self._client.stub.DictClear, req)
 
@@ -219,7 +220,7 @@ class _Dict(_Object, type_prefix="di"):
     async def __getitem__(self, key: Any) -> Any:
         """Get the value associated with a key.
 
-        This function only works in a synchronous context.
+        Note: this function will block the event loop when called in an async context.
         """
         NOT_FOUND = object()
         value = await self.get(key, NOT_FOUND)
@@ -233,7 +234,13 @@ class _Dict(_Object, type_prefix="di"):
         """Update the dictionary with additional items."""
         serialized = _serialize_dict(kwargs)
         req = api_pb2.DictUpdateRequest(dict_id=self.object_id, updates=serialized)
-        await retry_transient_errors(self._client.stub.DictUpdate, req)
+        try:
+            await retry_transient_errors(self._client.stub.DictUpdate, req)
+        except GRPCError as exc:
+            if "status = '413'" in exc.message:
+                raise RequestSizeError("Dict.update request is too large") from exc
+            else:
+                raise exc
 
     @live_method
     async def put(self, key: Any, value: Any) -> None:
@@ -241,13 +248,19 @@ class _Dict(_Object, type_prefix="di"):
         updates = {key: value}
         serialized = _serialize_dict(updates)
         req = api_pb2.DictUpdateRequest(dict_id=self.object_id, updates=serialized)
-        await retry_transient_errors(self._client.stub.DictUpdate, req)
+        try:
+            await retry_transient_errors(self._client.stub.DictUpdate, req)
+        except GRPCError as exc:
+            if "status = '413'" in exc.message:
+                raise RequestSizeError("Dict.put request is too large") from exc
+            else:
+                raise exc
 
     @live_method
     async def __setitem__(self, key: Any, value: Any) -> None:
         """Set a specific key-value pair to the dictionary.
 
-        This function only works in a synchronous context.
+        Note: this function will block the event loop when called in an async context.
         """
         return await self.put(key, value)
 
@@ -264,7 +277,7 @@ class _Dict(_Object, type_prefix="di"):
     async def __delitem__(self, key: Any) -> Any:
         """Delete a key from the dictionary.
 
-        This function only works in a synchronous context.
+        Note: this function will block the event loop when called in an async context.
         """
         return await self.pop(key)
 
@@ -272,9 +285,42 @@ class _Dict(_Object, type_prefix="di"):
     async def __contains__(self, key: Any) -> bool:
         """Return if a key is present.
 
-        This function only works in a synchronous context.
+        Note: this function will block the event loop when called in an async context.
         """
         return await self.contains(key)
 
+    @live_method_gen
+    async def keys(self) -> AsyncIterator[Any]:
+        """Return an iterator over the keys in this dictionary.
+
+        Note that (unlike with Python dicts) the return value is a simple iterator,
+        and results are unordered.
+        """
+        req = api_pb2.DictContentsRequest(dict_id=self.object_id, keys=True)
+        async for resp in self._client.stub.DictContents.unary_stream(req):
+            yield deserialize(resp.key, self._client)
+
+    @live_method_gen
+    async def values(self) -> AsyncIterator[Any]:
+        """Return an iterator over the values in this dictionary.
+
+        Note that (unlike with Python dicts) the return value is a simple iterator,
+        and results are unordered.
+        """
+        req = api_pb2.DictContentsRequest(dict_id=self.object_id, values=True)
+        async for resp in self._client.stub.DictContents.unary_stream(req):
+            yield deserialize(resp.value, self._client)
+
+    @live_method_gen
+    async def items(self) -> AsyncIterator[tuple[Any, Any]]:
+        """Return an iterator over the (key, value) tuples in this dictionary.
+
+        Note that (unlike with Python dicts) the return value is a simple iterator,
+        and results are unordered.
+        """
+        req = api_pb2.DictContentsRequest(dict_id=self.object_id, keys=True, values=True)
+        async for resp in self._client.stub.DictContents.unary_stream(req):
+            yield (deserialize(resp.key, self._client), deserialize(resp.value, self._client))
+
 
 Dict = synchronize_api(_Dict)