modal 0.66.14__py3-none-any.whl → 0.66.39__py3-none-any.whl

modal/experimental.py

@@ -7,7 +7,7 @@ from typing import (
 import modal._clustered_functions
 from modal.functions import _Function
 
-from ._container_io_manager import _ContainerIOManager
+from ._runtime.container_io_manager import _ContainerIOManager
 from .exception import (
     InvalidError,
 )
@@ -17,7 +17,6 @@ from .partial_function import _PartialFunction, _PartialFunctionFlags
 def stop_fetching_inputs():
     """Don't fetch any more inputs from the server, after the current one.
     The container will exit gracefully after the current input is processed."""
-
     _ContainerIOManager.stop_fetching_inputs()
 
 
@@ -25,7 +24,6 @@ def get_local_input_concurrency():
     """Get the container's local input concurrency.
     If recently reduced to particular value, it can return a larger number than
     set due to in-progress inputs."""
-
     return _ContainerIOManager.get_input_concurrency()
 
 
@@ -33,7 +31,6 @@ def set_local_input_concurrency(concurrency: int):
     """Set the container's local input concurrency. Dynamic concurrency will be disabled.
     When setting to a smaller value, this method will not interrupt in-progress inputs.
     """
-
     _ContainerIOManager.set_input_concurrency(concurrency)
 
 

modal/functions.py

@@ -35,6 +35,7 @@ from ._location import parse_cloud_provider
 from ._pty import get_pty_info
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
+from ._runtime.execution_context import current_input_id, is_local
 from ._serialization import serialize, serialize_proto_params
 from ._utils.async_utils import (
     TaskContext,
@@ -68,7 +69,6 @@ from .exception import (
     OutputExpiredError,
     deprecation_warning,
 )
-from .execution_context import current_input_id, is_local
 from .gpu import GPU_T, parse_gpu_config
 from .image import _Image
 from .mount import _get_client_mount, _Mount, get_auto_mounts
@@ -299,13 +299,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     """Functions are the basic units of serverless execution on Modal.
 
     Generally, you will not construct a `Function` directly. Instead, use the
-    `@app.function()` decorator on the `App` object (formerly called "Stub")
-    for your application.
+    `App.function()` decorator to register your Python functions with your App.
     """
 
     # TODO: more type annotations
     _info: Optional[FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[_Mount]  # set at load time, only by loader
+    _serve_mounts: typing.FrozenSet[_Mount]  # set at load time, only by loader
     _app: Optional["modal.app._App"] = None
     _obj: Optional["modal.cls._Obj"] = None  # only set for InstanceServiceFunctions and bound instance methods
     _web_url: Optional[str]
@@ -580,6 +579,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
         if is_local():
             entrypoint_mounts = info.get_entrypoint_mount()
+
             all_mounts = [
                 _get_client_mount(),
                 *explicit_mounts,
@@ -612,6 +612,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         if proxy:
             # HACK: remove this once we stop using ssh tunnels for this.
             if image:
+                # TODO(elias): this will cause an error if users use prior `.add_local_*` commands without copy=True
                 image = image.apt_install("autossh")
 
         function_spec = _FunctionSpec(
@@ -828,7 +829,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     )
                     for path, volume in validated_volumes
                 ]
-                loaded_mount_ids = {m.object_id for m in all_mounts}
+                loaded_mount_ids = {m.object_id for m in all_mounts} | {m.object_id for m in image._mount_layers}
 
                 # Get object dependencies
                 object_dependencies = []
@@ -970,9 +971,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         raise InvalidError(f"Function {info.function_name} is too large to deploy.")
                     raise
                 function_creation_status.set_response(response)
-            local_mounts = set(m for m in all_mounts if m.is_local())  # needed for modal.serve file watching
-            local_mounts |= image._used_local_mounts
-            obj._used_local_mounts = frozenset(local_mounts)
+            serve_mounts = set(m for m in all_mounts if m.is_local())  # needed for modal.serve file watching
+            serve_mounts |= image._serve_mounts
+            obj._serve_mounts = frozenset(serve_mounts)
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
         rep = f"Function({tag})"
@@ -1028,7 +1029,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 identity = "class service function for a parameterized class"
             if not self._parent.is_hydrated:
                 if self._parent.app._running_app is None:
-                    reason = ", because the App it is defined on is not running."
+                    reason = ", because the App it is defined on is not running"
                 else:
                     reason = ""
                 raise ExecutionError(
@@ -1223,7 +1224,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         self._function_name = None
         self._info = None
         self._use_function_id = ""
-        self._used_local_mounts = frozenset()
+        self._serve_mounts = frozenset()
 
     def _hydrate_metadata(self, metadata: Optional[Message]):
         # Overridden concrete implementation of base class method

modal/functions.pyi

@@ -110,7 +110,7 @@ OriginalReturnType = typing.TypeVar("OriginalReturnType", covariant=True)
 
 class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object._Object):
     _info: typing.Optional[modal._utils.function_utils.FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[modal.mount._Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount._Mount]
     _app: typing.Optional[modal.app._App]
     _obj: typing.Optional[modal.cls._Obj]
     _web_url: typing.Optional[str]
@@ -286,7 +286,7 @@ P_INNER = typing_extensions.ParamSpec("P_INNER")
 
 class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.Object):
     _info: typing.Optional[modal._utils.function_utils.FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[modal.mount.Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount.Mount]
     _app: typing.Optional[modal.app.App]
     _obj: typing.Optional[modal.cls.Obj]
     _web_url: typing.Optional[str]
@@ -450,11 +450,11 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
 
     _call_generator_nowait: ___call_generator_nowait_spec
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
 
-    remote: __remote_spec[P, ReturnType]
+    remote: __remote_spec[ReturnType, P]
 
     class __remote_gen_spec(typing_extensions.Protocol):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -466,17 +466,17 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     def _get_obj(self) -> typing.Optional[modal.cls.Obj]: ...
     def local(self, *args: P.args, **kwargs: P.kwargs) -> OriginalReturnType: ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    _experimental_spawn: ___experimental_spawn_spec[P, ReturnType]
+    _experimental_spawn: ___experimental_spawn_spec[ReturnType, P]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
         def __call__(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    spawn: __spawn_spec[P, ReturnType]
+    spawn: __spawn_spec[ReturnType, P]
 
     def get_raw_f(self) -> typing.Callable[..., typing.Any]: ...
 

modal/gpu.py

@@ -1,6 +1,6 @@
 # Copyright Modal Labs 2022
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from modal_proto import api_pb2
 
@@ -147,25 +147,27 @@ class Any(_GPUConfig):
         return f"GPU(Any, count={self.count})"
 
 
-STRING_TO_GPU_CONFIG = {
+STRING_TO_GPU_CONFIG: dict[str, Callable] = {
     "t4": T4,
     "l4": L4,
     "a100": A100,
+    "a100-80gb": lambda: A100(size="80GB"),
     "h100": H100,
     "a10g": A10G,
     "any": Any,
 }
-display_string_to_config = "\n".join(
-    f'- "{key}" → `{cls()}`' for key, cls in STRING_TO_GPU_CONFIG.items() if key != "inf2"
-)
+display_string_to_config = "\n".join(f'- "{key}" → `{c()}`' for key, c in STRING_TO_GPU_CONFIG.items() if key != "inf2")
 __doc__ = f"""
 **GPU configuration shortcodes**
 
 The following are the valid `str` values for the `gpu` parameter of
-[`@app.function`](/docs/reference/modal.Stub#function).
+[`@app.function`](/docs/reference/modal.App#function).
 
 {display_string_to_config}
 
+The shortcodes also support specifying count by suffixing `:N` to acquire `N` GPUs.
+For example, `a10g:4` will provision 4 A10G GPUs.
+
 Other configurations can be created using the constructors documented below.
 """
 

modal/image.py

@@ -273,12 +273,24 @@ class _Image(_Object, type_prefix="im"):
 
     force_build: bool
     inside_exceptions: List[Exception]
-    _used_local_mounts: typing.FrozenSet[_Mount]  # used for mounts watching
+    _serve_mounts: typing.FrozenSet[_Mount]  # used for mounts watching in `modal serve`
+    _deferred_mounts: Sequence[
+        _Mount
+    ]  # added as mounts on any container referencing the Image, see `def _mount_layers`
     _metadata: Optional[api_pb2.ImageMetadata] = None  # set on hydration, private for now
 
     def _initialize_from_empty(self):
         self.inside_exceptions = []
-        self._used_local_mounts = frozenset()
+        self._serve_mounts = frozenset()
+        self._deferred_mounts = ()
+        self.force_build = False
+
+    def _initialize_from_other(self, other: "_Image"):
+        # used by .clone()
+        self.inside_exceptions = other.inside_exceptions
+        self.force_build = other.force_build
+        self._serve_mounts = other._serve_mounts
+        self._deferred_mounts = other._deferred_mounts
 
     def _hydrate_metadata(self, message: Optional[Message]):
         env_image_id = config.get("image_id")  # set as an env var in containers
@@ -292,6 +304,51 @@ class _Image(_Object, type_prefix="im"):
             assert isinstance(message, api_pb2.ImageMetadata)
             self._metadata = message
 
+    def _add_mount_layer_or_copy(self, mount: _Mount, copy: bool = False):
+        if copy:
+            return self.copy_mount(mount, remote_path="/")
+
+        base_image = self
+
+        async def _load(self2: "_Image", resolver: Resolver, existing_object_id: Optional[str]):
+            self2._hydrate_from_other(base_image)  # same image id as base image as long as it's lazy
+            self2._deferred_mounts = tuple(base_image._deferred_mounts) + (mount,)
+            self2._serve_mounts = base_image._serve_mounts | ({mount} if mount.is_local() else set())
+
+        return _Image._from_loader(_load, "Image(local files)", deps=lambda: [base_image, mount])
+
+    @property
+    def _mount_layers(self) -> typing.Tuple[_Mount]:
+        """Non-evaluated mount layers on the image
+
+        When the image is used by a Modal container, these mounts need to be attached as well to
+        represent the full image content, as they haven't yet been represented as a layer in the
+        image.
+
+        When the image is used as a base image for a new layer (that is not itself a mount layer)
+        these mounts need to first be inserted as a copy operation (.copy_mount) into the image.
+        """
+        return self._deferred_mounts
+
+    def _assert_no_mount_layers(self):
+        if self._mount_layers:
+            raise InvalidError(
+                "An image tried to run a build step after using `image.add_local_*` to include local files.\n"
+                "\n"
+                "Run `image.add_local_*` commands last in your image build to avoid rebuilding images with every local "
+                "file change. Modal will then add these files to containers on startup instead, saving build time.\n"
+                "If you need to run other build steps after adding local files, set `copy=True` to copy the files"
+                "directly into the image, at the expense of some added build time.\n"
+                "\n"
+                "Example:\n"
+                "\n"
+                "my_image = (\n"
+                "    Image.debian_slim()\n"
+                '   .add_local_python_packages("mypak", copy=True)\n'
+                '    .run_commands("python -m mypak")  # this now works!\n'
+                ")\n"
+            )
+
     @staticmethod
     def _from_args(
         *,
@@ -306,9 +363,11 @@ class _Image(_Object, type_prefix="im"):
         force_build: bool = False,
         # For internal use only.
         _namespace: int = api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
+        _do_assert_no_mount_layers: bool = True,
     ):
         if base_images is None:
             base_images = {}
+
         if secrets is None:
             secrets = []
         if gpu_config is None:
@@ -334,6 +393,11 @@ class _Image(_Object, type_prefix="im"):
             return deps
 
         async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
+            if _do_assert_no_mount_layers:
+                for image in base_images.values():
+                    # base images can't have
+                    image._assert_no_mount_layers()
+
             environment = await _get_environment_cached(resolver.environment_name or "", resolver.client)
             # A bit hacky,but assume that the environment provides a valid builder version
             image_builder_version = cast(ImageBuilderVersion, environment._settings.image_builder_version)
@@ -349,7 +413,9 @@ class _Image(_Object, type_prefix="im"):
                     "No commands were provided for the image — have you tried using modal.Image.debian_slim()?"
                 )
             if dockerfile.commands and build_function:
-                raise InvalidError("Cannot provide both a build function and Dockerfile commands!")
+                raise InvalidError(
+                    "Cannot provide both build function and Dockerfile commands in the same image layer!"
+                )
 
             base_images_pb2s = [
                 api_pb2.BaseImage(
@@ -482,12 +548,12 @@ class _Image(_Object, type_prefix="im"):
             self._hydrate(image_id, resolver.client, result_response.metadata)
             local_mounts = set()
             for base in base_images.values():
-                local_mounts |= base._used_local_mounts
+                local_mounts |= base._serve_mounts
             if context_mount and context_mount.is_local():
                 local_mounts.add(context_mount)
-            self._used_local_mounts = frozenset(local_mounts)
+            self._serve_mounts = frozenset(local_mounts)
 
-        rep = "Image()"
+        rep = f"Image({dockerfile_function})"
         obj = _Image._from_loader(_load, rep, deps=_deps)
         obj.force_build = force_build
         return obj
@@ -553,6 +619,27 @@ class _Image(_Object, type_prefix="im"):
             context_mount=mount,
         )
 
+    def _add_local_python_packages(self, *packages: Union[str, Path], copy: bool = False) -> "_Image":
+        """Adds Python package files to containers
+
+        Adds all files from the specified Python packages to containers running the Image.
+
+        Packages are added to the `/root` directory of containers, which is on the `PYTHONPATH`
+        of any executed Modal Functions.
+
+        By default (`copy=False`), the files are added to containers on startup and are not built into the actual Image,
+        which speeds up deployment.
+
+        Set `copy=True` to copy the files into an Image layer at build time instead. This can slow down iteration since
+        it requires a rebuild of the Image and any subsequent build steps whenever the included files change, but it is
+        required if you want to run additional build steps after this one.
+
+        **Note:** This excludes all dot-prefixed subdirectories or files and all `.pyc`/`__pycache__` files.
+        To add full directories with finer control, use `.add_local_dir()` instead.
+        """
+        mount = _Mount.from_local_python_packages(*packages)
+        return self._add_mount_layer_or_copy(mount, copy=copy)
+
     def copy_local_dir(self, local_path: Union[str, Path], remote_path: Union[str, Path] = ".") -> "_Image":
         """Copy a directory into the image as a part of building the image.
 

modal/image.pyi

@@ -58,11 +58,17 @@ class DockerfileSpec:
 class _Image(modal.object._Object):
     force_build: bool
     inside_exceptions: typing.List[Exception]
-    _used_local_mounts: typing.FrozenSet[modal.mount._Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount._Mount]
+    _deferred_mounts: typing.Sequence[modal.mount._Mount]
     _metadata: typing.Optional[modal_proto.api_pb2.ImageMetadata]
 
     def _initialize_from_empty(self): ...
+    def _initialize_from_other(self, other: _Image): ...
     def _hydrate_metadata(self, message: typing.Optional[google.protobuf.message.Message]): ...
+    def _add_mount_layer_or_copy(self, mount: modal.mount._Mount, copy: bool = False): ...
+    @property
+    def _mount_layers(self) -> typing.Tuple[modal.mount._Mount]: ...
+    def _assert_no_mount_layers(self): ...
     @staticmethod
     def _from_args(
         *,
@@ -78,6 +84,7 @@ class _Image(modal.object._Object):
         context_mount: typing.Optional[modal.mount._Mount] = None,
         force_build: bool = False,
         _namespace: int = 1,
+        _do_assert_no_mount_layers: bool = True,
     ): ...
     def extend(
         self,
@@ -90,11 +97,13 @@ class _Image(modal.object._Object):
         context_mount: typing.Optional[modal.mount._Mount] = None,
         force_build: bool = False,
         _namespace: int = 1,
+        _do_assert_no_mount_layers: bool = True,
     ) -> _Image: ...
     def copy_mount(self, mount: modal.mount._Mount, remote_path: typing.Union[str, pathlib.Path] = ".") -> _Image: ...
     def copy_local_file(
         self, local_path: typing.Union[str, pathlib.Path], remote_path: typing.Union[str, pathlib.Path] = "./"
     ) -> _Image: ...
+    def _add_local_python_packages(self, *module_names, copy: bool = False) -> _Image: ...
     def copy_local_dir(
         self, local_path: typing.Union[str, pathlib.Path], remote_path: typing.Union[str, pathlib.Path] = "."
     ) -> _Image: ...
@@ -299,12 +308,18 @@ class _Image(modal.object._Object):
 class Image(modal.object.Object):
     force_build: bool
     inside_exceptions: typing.List[Exception]
-    _used_local_mounts: typing.FrozenSet[modal.mount.Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount.Mount]
+    _deferred_mounts: typing.Sequence[modal.mount.Mount]
     _metadata: typing.Optional[modal_proto.api_pb2.ImageMetadata]
 
     def __init__(self, *args, **kwargs): ...
     def _initialize_from_empty(self): ...
+    def _initialize_from_other(self, other: Image): ...
     def _hydrate_metadata(self, message: typing.Optional[google.protobuf.message.Message]): ...
+    def _add_mount_layer_or_copy(self, mount: modal.mount.Mount, copy: bool = False): ...
+    @property
+    def _mount_layers(self) -> typing.Tuple[modal.mount.Mount]: ...
+    def _assert_no_mount_layers(self): ...
     @staticmethod
     def _from_args(
         *,
@@ -320,6 +335,7 @@ class Image(modal.object.Object):
         context_mount: typing.Optional[modal.mount.Mount] = None,
         force_build: bool = False,
         _namespace: int = 1,
+        _do_assert_no_mount_layers: bool = True,
     ): ...
     def extend(
         self,
@@ -332,11 +348,13 @@ class Image(modal.object.Object):
         context_mount: typing.Optional[modal.mount.Mount] = None,
         force_build: bool = False,
         _namespace: int = 1,
+        _do_assert_no_mount_layers: bool = True,
     ) -> Image: ...
     def copy_mount(self, mount: modal.mount.Mount, remote_path: typing.Union[str, pathlib.Path] = ".") -> Image: ...
     def copy_local_file(
         self, local_path: typing.Union[str, pathlib.Path], remote_path: typing.Union[str, pathlib.Path] = "./"
     ) -> Image: ...
+    def _add_local_python_packages(self, *module_names, copy: bool = False) -> Image: ...
     def copy_local_dir(
         self, local_path: typing.Union[str, pathlib.Path], remote_path: typing.Union[str, pathlib.Path] = "."
     ) -> Image: ...

modal/io_streams.py

@@ -1,6 +1,18 @@
 # Copyright Modal Labs 2022
 import asyncio
-from typing import TYPE_CHECKING, AsyncGenerator, Generic, List, Literal, Optional, Tuple, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    AsyncGenerator,
+    AsyncIterator,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from grpclib import Status
 from grpclib.exceptions import GRPCError, StreamTerminatedError
@@ -18,7 +30,7 @@ if TYPE_CHECKING:
 
 
 async def _sandbox_logs_iterator(
-    sandbox_id: str, file_descriptor: int, last_entry_id: Optional[str], client: _Client
+    sandbox_id: str, file_descriptor: "api_pb2.FileDescriptor.ValueType", last_entry_id: str, client: _Client
 ) -> AsyncGenerator[Tuple[Optional[bytes], str], None]:
     req = api_pb2.SandboxGetLogsRequest(
         sandbox_id=sandbox_id,
@@ -37,7 +49,7 @@ async def _sandbox_logs_iterator(
 
 
 async def _container_process_logs_iterator(
-    process_id: str, file_descriptor: int, client: _Client
+    process_id: str, file_descriptor: "api_pb2.FileDescriptor.ValueType", client: _Client
 ) -> AsyncGenerator[Optional[bytes], None]:
     req = api_pb2.ContainerExecGetOutputRequest(
         exec_id=process_id, timeout=55, file_descriptor=file_descriptor, get_raw_bytes=True
@@ -74,9 +86,11 @@ class _StreamReader(Generic[T]):
     ```
     """
 
+    _stream: Optional[AsyncGenerator[Optional[bytes], None]]
+
     def __init__(
         self,
-        file_descriptor: int,
+        file_descriptor: "api_pb2.FileDescriptor.ValueType",
         object_id: str,
         object_type: Literal["sandbox", "container_process"],
         client: _Client,
@@ -90,7 +104,7 @@ class _StreamReader(Generic[T]):
         self._object_id = object_id
         self._client = client
         self._stream = None
-        self._last_entry_id: Optional[str] = None
+        self._last_entry_id: str = ""
         self._line_buffer = b""
 
         # Sandbox logs are streamed to the client as strings, so StreamReaders reading
@@ -145,15 +159,20 @@ class _StreamReader(Generic[T]):
         ```
 
         """
-        data = "" if self._text else b""
+        data_str = ""
+        data_bytes = b""
         async for message in self._get_logs():
             if message is None:
                 break
             if self._text:
-                data += message.decode("utf-8")
+                data_str += message.decode("utf-8")
             else:
-                data += message
-        return data
+                data_bytes += message
+
+        if self._text:
+            return cast(T, data_str)
+        else:
+            return cast(T, data_bytes)
 
     async def _consume_container_process_stream(self):
         """
@@ -275,7 +294,7 @@ class _StreamReader(Generic[T]):
                     line, self._line_buffer = self._line_buffer.split(b"\n", 1)
                     yield line + b"\n"
 
-    def __aiter__(self) -> AsyncGenerator[T, None]:
+    def __aiter__(self) -> AsyncIterator[T]:
         """mdmd:hidden"""
         if not self._stream:
             if self._by_line:
@@ -287,6 +306,7 @@ class _StreamReader(Generic[T]):
     async def __anext__(self) -> T:
         """mdmd:hidden"""
         assert self._stream is not None
+
         value = await self._stream.__anext__()
 
         # The stream yields None if it receives an EOF batch.
@@ -294,9 +314,9 @@ class _StreamReader(Generic[T]):
             raise StopAsyncIteration
 
         if self._text:
-            return value.decode("utf-8")
+            return cast(T, value.decode("utf-8"))
         else:
-            return value
+            return cast(T, value)
 
 
 MAX_BUFFER_SIZE = 2 * 1024 * 1024

modal/io_streams.pyi

@@ -4,7 +4,7 @@ import typing
 import typing_extensions
 
 def _sandbox_logs_iterator(
-    sandbox_id: str, file_descriptor: int, last_entry_id: typing.Optional[str], client: modal.client._Client
+    sandbox_id: str, file_descriptor: int, last_entry_id: str, client: modal.client._Client
 ) -> typing.AsyncGenerator[typing.Tuple[typing.Optional[bytes], str], None]: ...
 def _container_process_logs_iterator(
     process_id: str, file_descriptor: int, client: modal.client._Client
@@ -13,6 +13,8 @@ def _container_process_logs_iterator(
 T = typing.TypeVar("T")
 
 class _StreamReader(typing.Generic[T]):
+    _stream: typing.Optional[typing.AsyncGenerator[typing.Optional[bytes], None]]
+
     def __init__(
         self,
         file_descriptor: int,
@@ -30,7 +32,7 @@ class _StreamReader(typing.Generic[T]):
     def _stream_container_process(self) -> typing.AsyncGenerator[typing.Tuple[typing.Optional[bytes], str], None]: ...
     def _get_logs(self) -> typing.AsyncGenerator[typing.Optional[bytes], None]: ...
     def _get_logs_by_line(self) -> typing.AsyncGenerator[typing.Optional[bytes], None]: ...
-    def __aiter__(self) -> typing.AsyncGenerator[T, None]: ...
+    def __aiter__(self) -> typing.AsyncIterator[T]: ...
     async def __anext__(self) -> T: ...
 
 class _StreamWriter:
@@ -45,6 +47,8 @@ class _StreamWriter:
 T_INNER = typing.TypeVar("T_INNER", covariant=True)
 
 class StreamReader(typing.Generic[T]):
+    _stream: typing.Optional[typing.AsyncGenerator[typing.Optional[bytes], None]]
+
     def __init__(
         self,
         file_descriptor: int,
@@ -88,8 +92,8 @@ class StreamReader(typing.Generic[T]):
 
     _get_logs_by_line: ___get_logs_by_line_spec
 
-    def __iter__(self) -> typing.Generator[T, None, None]: ...
-    def __aiter__(self) -> typing.AsyncGenerator[T, None]: ...
+    def __iter__(self) -> typing.Iterator[T]: ...
+    def __aiter__(self) -> typing.AsyncIterator[T]: ...
     def __next__(self) -> T: ...
     async def __anext__(self) -> T: ...
 

modal/mount.py

@@ -571,7 +571,7 @@ class _Mount(_Object, type_prefix="mo"):
         # Don't re-run inside container.
 
         mount = _Mount._new()
-        from .execution_context import is_local
+        from ._runtime.execution_context import is_local
 
         if not is_local():
             return mount  # empty/non-mountable mount in case it's used from within a container
@@ -622,12 +622,13 @@ class _Mount(_Object, type_prefix="mo"):
         client: Optional[_Client] = None,
     ) -> None:
         check_object_name(deployment_name, "Mount")
+        environment_name = _get_environment_name(environment_name, resolver=None)
         self._deployment_name = deployment_name
         self._namespace = namespace
         self._environment_name = environment_name
         if client is None:
             client = await _Client.from_env()
-        resolver = Resolver(client=client)
+        resolver = Resolver(client=client, environment_name=environment_name)
         await resolver.load(self)
 
     def _get_metadata(self) -> api_pb2.MountHandleMetadata:

modal/network_file_system.py

@@ -171,34 +171,6 @@ class _NetworkFileSystem(_Object, type_prefix="sv"):
             tc.infinite_loop(lambda: client.stub.SharedVolumeHeartbeat(request), sleep=_heartbeat_sleep)
             yield cls._new_hydrated(response.shared_volume_id, client, None, is_another_app=True)
 
-    @staticmethod
-    def persisted(
-        label: str,
-        namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
-        environment_name: Optional[str] = None,
-        cloud: Optional[str] = None,
-    ):
-        """mdmd:hidden"""
-        message = (
-            "`NetworkFileSystem.persisted` is deprecated."
-            " Please use `NetworkFileSystem.from_name(name, create_if_missing=True)` instead."
-        )
-        deprecation_error((2024, 3, 1), message)
-
-    def persist(
-        self,
-        label: str,
-        namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
-        environment_name: Optional[str] = None,
-        cloud: Optional[str] = None,
-    ):
-        """mdmd:hidden"""
-        message = (
-            "`NetworkFileSystem().persist('my-volume')` is deprecated."
-            " Please use `NetworkFileSystem.from_name('my-volume', create_if_missing=True)` instead."
-        )
-        deprecation_error((2024, 2, 29), message)
-
     @staticmethod
     async def lookup(
         label: str,