modal 0.70.2__py3-none-any.whl → 0.71.2__py3-none-any.whl

modal/_container_entrypoint.py

@@ -11,7 +11,6 @@ if telemetry_socket:
     instrument_imports(telemetry_socket)
 
 import asyncio
-import base64
 import concurrent.futures
 import inspect
 import queue
@@ -337,14 +336,17 @@ def call_function(
                     signal.signal(signal.SIGUSR1, usr1_handler)  # reset signal handler
 
 
-def get_active_app_fallback(function_def: api_pb2.Function) -> Optional[_App]:
+def get_active_app_fallback(function_def: api_pb2.Function) -> _App:
     # This branch is reached in the special case that the imported function/class is:
     # 1) not serialized, and
     # 2) isn't a FunctionHandle - i.e, not decorated at definition time
     # Look at all instantiated apps - if there is only one with the indicated name, use that one
     app_name: Optional[str] = function_def.app_name or None  # coalesce protobuf field to None
     matching_apps = _App._all_apps.get(app_name, [])
-    active_app = None
+    if len(matching_apps) == 1:
+        active_app: _App = matching_apps[0]
+        return active_app
+
     if len(matching_apps) > 1:
         if app_name is not None:
             warning_sub_message = f"app with the same name ('{app_name}')"
@@ -354,12 +356,10 @@ def get_active_app_fallback(function_def: api_pb2.Function) -> Optional[_App]:
             f"You have more than one {warning_sub_message}. "
             "It's recommended to name all your Apps uniquely when using multiple apps"
         )
-    elif len(matching_apps) == 1:
-        (active_app,) = matching_apps
-    # there could also technically be zero found apps, but that should probably never be an
-    # issue since that would mean user won't use is_inside or other function handles anyway
 
-    return active_app
+    # If we don't have an active app, create one on the fly
+    # The app object is used to carry the app layout etc
+    return _App()
 
 
 def call_lifecycle_functions(
@@ -403,7 +403,7 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
     # This is a bit weird but we need both the blocking and async versions of ContainerIOManager.
     # At some point, we should fix that by having built-in support for running "user code"
     container_io_manager = ContainerIOManager(container_args, client)
-    active_app: Optional[_App] = None
+    active_app: _App
     service: Service
     function_def = container_args.function_def
     is_auto_snapshot: bool = function_def.is_auto_snapshot
@@ -450,8 +450,9 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
                 )
 
             # If the cls/function decorator was applied in local scope, but the app is global, we can look it up
-            active_app = service.app
-            if active_app is None:
+            if service.app is not None:
+                active_app = service.app
+            else:
                 # if the app can't be inferred by the imported function, use name-based fallback
                 active_app = get_active_app_fallback(function_def)
 
@@ -468,9 +469,8 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
 
         # Initialize objects on the app.
         # This is basically only functions and classes - anything else is deprecated and will be unsupported soon
-        if active_app is not None:
-            app: App = synchronizer._translate_out(active_app)
-            app._init_container(client, container_app)
+        app: App = synchronizer._translate_out(active_app)
+        app._init_container(client, container_app)
 
         # Hydrate all function dependencies.
         # TODO(erikbern): we an remove this once we
@@ -581,7 +581,15 @@ if __name__ == "__main__":
     logger.debug("Container: starting")
 
     container_args = api_pb2.ContainerArguments()
-    container_args.ParseFromString(base64.b64decode(sys.argv[1]))
+
+    container_arguments_path: Optional[str] = os.environ.get("MODAL_CONTAINER_ARGUMENTS_PATH")
+    if container_arguments_path is None:
+        # TODO(erikbern): this fallback is for old workers and we can remove it very soon (days)
+        import base64
+
+        container_args.ParseFromString(base64.b64decode(sys.argv[1]))
+    else:
+        container_args.ParseFromString(open(container_arguments_path, "rb").read())
 
     # Note that we're creating the client in a synchronous context, but it will be running in a separate thread.
     # This is good because if the function is long running then we the client can still send heartbeats

modal/_runtime/container_io_manager.py

@@ -457,12 +457,7 @@ class _ContainerIOManager:
             resp = await retry_transient_errors(self._client.stub.AppGetLayout, req)
             app_layout = resp.app_layout
 
-        return running_app_from_layout(
-            self.app_id,
-            app_layout,
-            self._client,
-            environment_name=self._environment_name,
-        )
+        return running_app_from_layout(self.app_id, app_layout)
 
     async def get_serialized_function(self) -> tuple[Optional[Any], Optional[Callable[..., Any]]]:
         # Fetch the serialized function definition

modal/app.py

@@ -168,7 +168,7 @@ class _App:
     """
 
     _all_apps: ClassVar[dict[Optional[str], list["_App"]]] = {}
-    _container_app: ClassVar[Optional[RunningApp]] = None
+    _container_app: ClassVar[Optional["_App"]] = None
 
     _name: Optional[str]
     _description: Optional[str]
@@ -294,12 +294,7 @@ class _App:
         app = _App(name)
         app._app_id = response.app_id
         app._client = client
-        app._running_app = RunningApp(
-            response.app_id,
-            client=client,
-            environment_name=environment_name,
-            interactive=False,
-        )
+        app._running_app = RunningApp(response.app_id, interactive=False)
         return app
 
     def set_description(self, description: str):
@@ -463,8 +458,8 @@ class _App:
         if self._running_app:
             # If this is inside a container, then objects can be defined after app initialization.
             # So we may have to initialize objects once they get bound to the app.
-            if function.tag in self._running_app.tag_to_object_id:
-                object_id: str = self._running_app.tag_to_object_id[function.tag]
+            if function.tag in self._running_app.function_ids:
+                object_id: str = self._running_app.function_ids[function.tag]
                 metadata: Message = self._running_app.object_handle_metadata[object_id]
                 function._hydrate(object_id, self._client, metadata)
 
@@ -476,8 +471,8 @@ class _App:
         if self._running_app:
             # If this is inside a container, then objects can be defined after app initialization.
             # So we may have to initialize objects once they get bound to the app.
-            if tag in self._running_app.tag_to_object_id:
-                object_id: str = self._running_app.tag_to_object_id[tag]
+            if tag in self._running_app.class_ids:
+                object_id: str = self._running_app.class_ids[tag]
                 metadata: Message = self._running_app.object_handle_metadata[object_id]
                 cls._hydrate(object_id, self._client, metadata)
 
@@ -488,21 +483,21 @@ class _App:
         self._running_app = running_app
         self._client = client
 
-        _App._container_app = running_app
-
-        # Hydrate objects on app -- hydrating functions first so that when a class is being hydrated its
-        # corresponding class service function is already hydrated.
-        def hydrate_objects(objects_dict):
-            for tag, object_id in running_app.tag_to_object_id.items():
-                if tag in objects_dict:
-                    obj = objects_dict[tag]
-                    handle_metadata = running_app.object_handle_metadata[object_id]
-                    obj._hydrate(object_id, client, handle_metadata)
+        _App._container_app = self
 
         # Hydrate function objects
-        hydrate_objects(self._functions)
+        for tag, object_id in running_app.function_ids.items():
+            if tag in self._functions:
+                obj = self._functions[tag]
+                handle_metadata = running_app.object_handle_metadata[object_id]
+                obj._hydrate(object_id, client, handle_metadata)
+
         # Hydrate class objects
-        hydrate_objects(self._classes)
+        for tag, object_id in running_app.class_ids.items():
+            if tag in self._classes:
+                obj = self._classes[tag]
+                handle_metadata = running_app.object_handle_metadata[object_id]
+                obj._hydrate(object_id, client, handle_metadata)
 
     @property
     def registered_functions(self) -> dict[str, _Function]:
@@ -1047,6 +1042,13 @@ class _App:
                     if log.data:
                         yield log.data
 
+    @classmethod
+    def _get_container_app(cls) -> Optional["_App"]:
+        """Returns the `App` running inside a container.
+
+        This will return `None` outside of a Modal container."""
+        return cls._container_app
+
     @classmethod
     def _reset_container_app(cls):
         """Only used for tests."""

modal/app.pyi

@@ -73,7 +73,7 @@ class _FunctionDecoratorType:
 
 class _App:
     _all_apps: typing.ClassVar[dict[typing.Optional[str], list[_App]]]
-    _container_app: typing.ClassVar[typing.Optional[modal.running_app.RunningApp]]
+    _container_app: typing.ClassVar[typing.Optional[_App]]
     _name: typing.Optional[str]
     _description: typing.Optional[str]
     _functions: dict[str, modal.functions._Function]
@@ -266,11 +266,13 @@ class _App:
         self, client: typing.Optional[modal.client._Client] = None
     ) -> collections.abc.AsyncGenerator[str, None]: ...
     @classmethod
+    def _get_container_app(cls) -> typing.Optional[_App]: ...
+    @classmethod
     def _reset_container_app(cls): ...
 
 class App:
     _all_apps: typing.ClassVar[dict[typing.Optional[str], list[App]]]
-    _container_app: typing.ClassVar[typing.Optional[modal.running_app.RunningApp]]
+    _container_app: typing.ClassVar[typing.Optional[App]]
     _name: typing.Optional[str]
     _description: typing.Optional[str]
     _functions: dict[str, modal.functions.Function]
@@ -530,6 +532,8 @@ class App:
 
     _logs: ___logs_spec
 
+    @classmethod
+    def _get_container_app(cls) -> typing.Optional[App]: ...
     @classmethod
     def _reset_container_app(cls): ...
 

modal/client.pyi

@@ -26,7 +26,7 @@ class _Client:
     _stub: typing.Optional[modal_proto.api_grpc.ModalClientStub]
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.70.2"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.71.2"
     ): ...
     def is_closed(self) -> bool: ...
     @property
@@ -81,7 +81,7 @@ class Client:
     _stub: typing.Optional[modal_proto.api_grpc.ModalClientStub]
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.70.2"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.71.2"
     ): ...
     def is_closed(self) -> bool: ...
     @property

modal/experimental.py

@@ -48,6 +48,9 @@ def clustered(size: int, broadcast: bool = True):
 
     assert broadcast, "broadcast=False has not been implemented yet!"
 
+    if size <= 0:
+        raise ValueError("cluster size must be greater than 0")
+
     def wrapper(raw_f: Callable[..., Any]) -> _PartialFunction:
         if isinstance(raw_f, _Function):
             raw_f = raw_f.get_raw_f()

modal/file_io.py

@@ -1,6 +1,8 @@
 # Copyright Modal Labs 2024
 import asyncio
+import enum
 import io
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, AsyncIterator, Generic, Optional, Sequence, TypeVar, Union, cast
 
 if TYPE_CHECKING:
@@ -11,6 +13,7 @@ import json
 from grpclib.exceptions import GRPCError, StreamTerminatedError
 
 from modal._utils.grpc_utils import retry_transient_errors
+from modal.io_streams_helper import consume_stream_with_retries
 from modal_proto import api_pb2
 
 from ._utils.async_utils import synchronize_api
@@ -94,6 +97,20 @@ async def _replace_bytes(file: "_FileIO", data: bytes, start: Optional[int] = No
     await file._wait(resp.exec_id)
 
 
+class FileWatchEventType(enum.Enum):
+    Unknown = "Unknown"
+    Access = "Access"
+    Create = "Create"
+    Modify = "Modify"
+    Remove = "Remove"
+
+
+@dataclass
+class FileWatchEvent:
+    paths: list[str]
+    type: FileWatchEventType
+
+
 # The FileIO class is designed to mimic Python's io.FileIO
 # See https://github.com/python/cpython/blob/main/Lib/_pyio.py#L1459
 class _FileIO(Generic[T]):
@@ -124,6 +141,7 @@ class _FileIO(Generic[T]):
     _task_id: str = ""
     _file_descriptor: str = ""
     _client: Optional[_Client] = None
+    _watch_output_buffer: list[Optional[bytes]] = []
 
     def _validate_mode(self, mode: str) -> None:
         if not any(char in mode for char in "rwax"):
@@ -167,6 +185,44 @@ class _FileIO(Generic[T]):
             for message in batch.output:
                 yield message
 
+    async def _consume_watch_output(self, exec_id: str) -> None:
+        def item_handler(item: Optional[bytes]):
+            self._watch_output_buffer.append(item)
+
+        def completion_check(item: Optional[bytes]):
+            return item is None
+
+        await consume_stream_with_retries(
+            self._consume_output(exec_id),
+            item_handler,
+            completion_check,
+        )
+
+    async def _parse_watch_output(self, event: bytes) -> Optional[FileWatchEvent]:
+        try:
+            event_json = json.loads(event.decode())
+            return FileWatchEvent(type=FileWatchEventType(event_json["event_type"]), paths=event_json["paths"])
+        except (json.JSONDecodeError, KeyError, ValueError):
+            # skip invalid events
+            return None
+
+    async def _stream_watch_output(self) -> AsyncIterator[FileWatchEvent]:
+        buffer = b""
+        while True:
+            if len(self._watch_output_buffer) > 0:
+                item = self._watch_output_buffer.pop(0)
+                if item is None:
+                    break
+                buffer += item
+                # a single event may be split across multiple messages, the end of an event is marked by two newlines
+                if buffer.endswith(b"\n\n"):
+                    event = await self._parse_watch_output(buffer.strip())
+                    if event is not None:
+                        yield event
+                    buffer = b""
+            else:
+                await asyncio.sleep(0.1)
+
     async def _wait(self, exec_id: str) -> bytes:
         # The logic here is similar to how output is read from `exec`
         output = b""
@@ -391,6 +447,36 @@ class _FileIO(Generic[T]):
         )
         await self._wait(resp.exec_id)
 
+    @classmethod
+    async def watch(
+        cls,
+        path: str,
+        client: _Client,
+        task_id: str,
+        filter: Optional[list[FileWatchEventType]] = None,
+        recursive: bool = False,
+        timeout: Optional[int] = None,
+    ) -> AsyncIterator[FileWatchEvent]:
+        self = cls.__new__(cls)
+        self._client = client
+        self._task_id = task_id
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_watch_request=api_pb2.ContainerFileWatchRequest(
+                    path=path,
+                    recursive=recursive,
+                    timeout_secs=timeout,
+                ),
+                task_id=self._task_id,
+            )
+        )
+        task = asyncio.create_task(self._consume_watch_output(resp.exec_id))
+        async for event in self._stream_watch_output():
+            if filter and event.type not in filter:
+                continue
+            yield event
+        task.cancel()
+
     async def _close(self) -> None:
         # Buffer is flushed by the runner on close
         resp = await self._make_request(

modal/file_io.pyi

@@ -1,4 +1,5 @@
 import _typeshed
+import enum
 import modal.client
 import modal_proto.api_pb2
 import typing
@@ -13,14 +14,33 @@ async def _replace_bytes(
     file: _FileIO, data: bytes, start: typing.Optional[int] = None, end: typing.Optional[int] = None
 ) -> None: ...
 
+class FileWatchEventType(enum.Enum):
+    Unknown = "Unknown"
+    Access = "Access"
+    Create = "Create"
+    Modify = "Modify"
+    Remove = "Remove"
+
+class FileWatchEvent:
+    paths: list[str]
+    type: FileWatchEventType
+
+    def __init__(self, paths: list[str], type: FileWatchEventType) -> None: ...
+    def __repr__(self): ...
+    def __eq__(self, other): ...
+
 class _FileIO(typing.Generic[T]):
     _task_id: str
     _file_descriptor: str
     _client: typing.Optional[modal.client._Client]
+    _watch_output_buffer: list[typing.Optional[bytes]]
 
     def _validate_mode(self, mode: str) -> None: ...
     def _handle_error(self, error: modal_proto.api_pb2.SystemErrorMessage) -> None: ...
     def _consume_output(self, exec_id: str) -> typing.AsyncIterator[typing.Optional[bytes]]: ...
+    async def _consume_watch_output(self, exec_id: str) -> None: ...
+    async def _parse_watch_output(self, event: bytes) -> typing.Optional[FileWatchEvent]: ...
+    def _stream_watch_output(self) -> typing.AsyncIterator[FileWatchEvent]: ...
     async def _wait(self, exec_id: str) -> bytes: ...
     def _validate_type(self, data: typing.Union[bytes, str]) -> None: ...
     async def _open_file(self, path: str, mode: str) -> None: ...
@@ -49,6 +69,16 @@ class _FileIO(typing.Generic[T]):
     async def mkdir(cls, path: str, client: modal.client._Client, task_id: str, parents: bool = False) -> None: ...
     @classmethod
     async def rm(cls, path: str, client: modal.client._Client, task_id: str, recursive: bool = False) -> None: ...
+    @classmethod
+    def watch(
+        cls,
+        path: str,
+        client: modal.client._Client,
+        task_id: str,
+        filter: typing.Optional[list[FileWatchEventType]] = None,
+        recursive: bool = False,
+        timeout: typing.Optional[int] = None,
+    ) -> typing.AsyncIterator[FileWatchEvent]: ...
     async def _close(self) -> None: ...
     async def close(self) -> None: ...
     def _check_writable(self) -> None: ...
@@ -79,6 +109,7 @@ class FileIO(typing.Generic[T]):
     _task_id: str
     _file_descriptor: str
     _client: typing.Optional[modal.client.Client]
+    _watch_output_buffer: list[typing.Optional[bytes]]
 
     def __init__(self, /, *args, **kwargs): ...
     def _validate_mode(self, mode: str) -> None: ...
@@ -90,6 +121,24 @@ class FileIO(typing.Generic[T]):
 
     _consume_output: ___consume_output_spec
 
+    class ___consume_watch_output_spec(typing_extensions.Protocol):
+        def __call__(self, exec_id: str) -> None: ...
+        async def aio(self, exec_id: str) -> None: ...
+
+    _consume_watch_output: ___consume_watch_output_spec
+
+    class ___parse_watch_output_spec(typing_extensions.Protocol):
+        def __call__(self, event: bytes) -> typing.Optional[FileWatchEvent]: ...
+        async def aio(self, event: bytes) -> typing.Optional[FileWatchEvent]: ...
+
+    _parse_watch_output: ___parse_watch_output_spec
+
+    class ___stream_watch_output_spec(typing_extensions.Protocol):
+        def __call__(self) -> typing.Iterator[FileWatchEvent]: ...
+        def aio(self) -> typing.AsyncIterator[FileWatchEvent]: ...
+
+    _stream_watch_output: ___stream_watch_output_spec
+
     class ___wait_spec(typing_extensions.Protocol):
         def __call__(self, exec_id: str) -> bytes: ...
         async def aio(self, exec_id: str) -> bytes: ...
@@ -173,6 +222,16 @@ class FileIO(typing.Generic[T]):
     def mkdir(cls, path: str, client: modal.client.Client, task_id: str, parents: bool = False) -> None: ...
     @classmethod
     def rm(cls, path: str, client: modal.client.Client, task_id: str, recursive: bool = False) -> None: ...
+    @classmethod
+    def watch(
+        cls,
+        path: str,
+        client: modal.client.Client,
+        task_id: str,
+        filter: typing.Optional[list[FileWatchEventType]] = None,
+        recursive: bool = False,
+        timeout: typing.Optional[int] = None,
+    ) -> typing.Iterator[FileWatchEvent]: ...
 
     class ___close_spec(typing_extensions.Protocol):
         def __call__(self) -> None: ...

modal/functions.pyi

@@ -462,11 +462,11 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
 
     _call_generator_nowait: ___call_generator_nowait_spec
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_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[ReturnType, P]
+    remote: __remote_spec[P, ReturnType]
 
     class __remote_gen_spec(typing_extensions.Protocol):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -479,17 +479,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[ReturnType_INNER, P_INNER]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_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[ReturnType, P]
+    _experimental_spawn: ___experimental_spawn_spec[P, ReturnType]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_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[ReturnType, P]
+    spawn: __spawn_spec[P, ReturnType]
 
     def get_raw_f(self) -> typing.Callable[..., typing.Any]: ...
 

modal/io_streams.py

@@ -14,7 +14,8 @@ from typing import (
 from grpclib import Status
 from grpclib.exceptions import GRPCError, StreamTerminatedError
 
-from modal.exception import ClientClosed, InvalidError
+from modal.exception import InvalidError
+from modal.io_streams_helper import consume_stream_with_retries
 from modal_proto import api_pb2
 
 from ._utils.async_utils import synchronize_api
@@ -176,34 +177,21 @@ class _StreamReader(Generic[T]):
         if self._stream_type == StreamType.DEVNULL:
             return
 
-        completed = False
-        retries_remaining = 10
-        while not completed:
-            try:
-                iterator = _container_process_logs_iterator(self._object_id, self._file_descriptor, self._client)
+        def item_handler(item: Optional[bytes]):
+            if self._stream_type == StreamType.STDOUT and item is not None:
+                print(item.decode("utf-8"), end="")
+            elif self._stream_type == StreamType.PIPE:
+                self._container_process_buffer.append(item)
 
-                async for message in iterator:
-                    if self._stream_type == StreamType.STDOUT and message:
-                        print(message.decode("utf-8"), end="")
-                    elif self._stream_type == StreamType.PIPE:
-                        self._container_process_buffer.append(message)
-                    if message is None:
-                        completed = True
-                        break
+        def completion_check(item: Optional[bytes]):
+            return item is None
 
-            except (GRPCError, StreamTerminatedError, ClientClosed) as exc:
-                if retries_remaining > 0:
-                    retries_remaining -= 1
-                    if isinstance(exc, GRPCError):
-                        if exc.status in RETRYABLE_GRPC_STATUS_CODES:
-                            await asyncio.sleep(1.0)
-                            continue
-                    elif isinstance(exc, StreamTerminatedError):
-                        continue
-                    elif isinstance(exc, ClientClosed):
-                        # If the client was closed, the user has triggered a cleanup.
-                        break
-                raise exc
+        iterator = _container_process_logs_iterator(self._object_id, self._file_descriptor, self._client)
+        await consume_stream_with_retries(
+            iterator,
+            item_handler,
+            completion_check,
+        )
 
     async def _stream_container_process(self) -> AsyncGenerator[tuple[Optional[bytes], str], None]:
         """Streams the container process buffer to the reader."""

modal/io_streams_helper.py

@@ -0,0 +1,53 @@
+# Copyright Modal Labs 2024
+import asyncio
+from typing import AsyncIterator, Callable, TypeVar
+
+from grpclib.exceptions import GRPCError, StreamTerminatedError
+
+from modal.exception import ClientClosed
+
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES
+
+T = TypeVar("T")
+
+
+async def consume_stream_with_retries(
+    stream: AsyncIterator[T],
+    item_handler: Callable[[T], None],
+    completion_check: Callable[[T], bool],
+    max_retries: int = 10,
+    retry_delay: float = 1.0,
+) -> None:
+    """mdmd:hidden
+    Helper function to consume a stream with retry logic for transient errors.
+
+    Args:
+        stream_generator: Function that returns an AsyncIterator to consume
+        item_handler: Callback function to handle each item from the stream
+        completion_check: Callback function to check if the stream is complete
+        max_retries: Maximum number of retry attempts
+        retry_delay: Delay in seconds between retries
+    """
+    completed = False
+    retries_remaining = max_retries
+
+    while not completed:
+        try:
+            async for item in stream:
+                item_handler(item)
+                if completion_check(item):
+                    completed = True
+                    break
+
+        except (GRPCError, StreamTerminatedError, ClientClosed) as exc:
+            if retries_remaining > 0:
+                retries_remaining -= 1
+                if isinstance(exc, GRPCError):
+                    if exc.status in RETRYABLE_GRPC_STATUS_CODES:
+                        await asyncio.sleep(retry_delay)
+                        continue
+                elif isinstance(exc, StreamTerminatedError):
+                    continue
+                elif isinstance(exc, ClientClosed):
+                    break
+            raise