modal 1.2.1.dev19__py3-none-any.whl → 1.2.2.dev19__py3-none-any.whl

modal/image.py

@@ -38,7 +38,8 @@ from ._utils.docker_utils import (
     find_dockerignore_file,
 )
 from ._utils.function_utils import FunctionInfo
-from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES
+from ._utils.mount_utils import validate_only_modal_volumes
 from .client import _Client
 from .cloud_bucket_mount import _CloudBucketMount
 from .config import config, logger, user_config_path
@@ -487,6 +488,7 @@ class _Image(_Object, type_prefix="im"):
         context_mount_function: Optional[Callable[[], Optional[_Mount]]] = None,
         force_build: bool = False,
         build_args: dict[str, str] = {},
+        validated_volumes: Optional[Sequence[tuple[str, _Volume]]] = None,
         # For internal use only.
         _namespace: "api_pb2.DeploymentNamespace.ValueType" = api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
         _do_assert_no_mount_layers: bool = True,
@@ -494,6 +496,9 @@ class _Image(_Object, type_prefix="im"):
         if base_images is None:
             base_images = {}
 
+        if validated_volumes is None:
+            validated_volumes = []
+
         if secrets is None:
             secrets = []
         if gpu_config is None:
@@ -514,6 +519,8 @@ class _Image(_Object, type_prefix="im"):
                 deps += (build_function,)
             if image_registry_config and image_registry_config.secret:
                 deps += (image_registry_config.secret,)
+            for _, vol in validated_volumes:
+                deps += (vol,)
             return deps
 
         async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
@@ -592,6 +599,17 @@ class _Image(_Object, type_prefix="im"):
                 build_function_id = ""
                 _build_function = None
 
+            # Relies on dicts being ordered (true as of Python 3.6).
+            volume_mounts = [
+                api_pb2.VolumeMount(
+                    mount_path=path,
+                    volume_id=volume.object_id,
+                    allow_background_commits=True,
+                    read_only=volume._read_only,
+                )
+                for path, volume in validated_volumes
+            ]
+
             image_definition = api_pb2.Image(
                 base_images=base_images_pb2s,
                 dockerfile_commands=dockerfile.commands,
@@ -604,6 +622,7 @@ class _Image(_Object, type_prefix="im"):
                 runtime_debug=config.get("function_runtime_debug"),
                 build_function=_build_function,
                 build_args=build_args,
+                volume_mounts=volume_mounts,
             )
 
             req = api_pb2.ImageGetOrCreateRequest(
@@ -619,7 +638,7 @@ class _Image(_Object, type_prefix="im"):
                 allow_global_deployment=os.environ.get("MODAL_IMAGE_ALLOW_GLOBAL_DEPLOYMENT") == "1",
                 ignore_cache=config.get("ignore_cache"),
             )
-            resp = await retry_transient_errors(resolver.client.stub.ImageGetOrCreate, req)
+            resp = await resolver.client.stub.ImageGetOrCreate(req)
             image_id = resp.image_id
             result: api_pb2.GenericResult
             metadata: Optional[api_pb2.ImageMetadata] = None
@@ -848,7 +867,7 @@ class _Image(_Object, type_prefix="im"):
             client = await _Client.from_env()
 
         async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
-            resp = await retry_transient_errors(client.stub.ImageFromId, api_pb2.ImageFromIdRequest(image_id=image_id))
+            resp = await client.stub.ImageFromId(api_pb2.ImageFromIdRequest(image_id=image_id))
             self._hydrate(resp.image_id, resolver.client, resp.metadata)
 
         rep = f"Image.from_id({image_id!r})"
@@ -1690,6 +1709,7 @@ class _Image(_Object, type_prefix="im"):
         *commands: Union[str, list[str]],
         env: Optional[dict[str, Optional[str]]] = None,
         secrets: Optional[Collection[_Secret]] = None,
+        volumes: Optional[dict[Union[str, PurePosixPath], _Volume]] = None,
         gpu: GPU_T = None,
         force_build: bool = False,  # Ignore cached builds, similar to 'docker build --no-cache'
     ) -> "_Image":
@@ -1712,6 +1732,7 @@ class _Image(_Object, type_prefix="im"):
             secrets=secrets,
             gpu_config=parse_gpu_config(gpu),
             force_build=self.force_build or force_build,
+            validated_volumes=validate_only_modal_volumes(volumes, "Image.run_commands"),
         )
 
     @staticmethod

modal/image.pyi

@@ -176,6 +176,7 @@ class _Image(modal._object._Object):
         ] = None,
         force_build: bool = False,
         build_args: dict[str, str] = {},
+        validated_volumes: typing.Optional[collections.abc.Sequence[tuple[str, modal.volume._Volume]]] = None,
         _namespace: int = 1,
         _do_assert_no_mount_layers: bool = True,
     ): ...
@@ -668,6 +669,7 @@ class _Image(modal._object._Object):
         *commands: typing.Union[str, list[str]],
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret._Secret]] = None,
+        volumes: typing.Optional[dict[typing.Union[str, pathlib.PurePosixPath], modal.volume._Volume]] = None,
         gpu: typing.Union[None, str, modal.gpu._GPUConfig] = None,
         force_build: bool = False,
     ) -> _Image:
@@ -1091,6 +1093,7 @@ class Image(modal.object.Object):
         ] = None,
         force_build: bool = False,
         build_args: dict[str, str] = {},
+        validated_volumes: typing.Optional[collections.abc.Sequence[tuple[str, modal.volume.Volume]]] = None,
         _namespace: int = 1,
         _do_assert_no_mount_layers: bool = True,
     ): ...
@@ -1648,6 +1651,7 @@ class Image(modal.object.Object):
         *commands: typing.Union[str, list[str]],
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret.Secret]] = None,
+        volumes: typing.Optional[dict[typing.Union[str, pathlib.PurePosixPath], modal.volume.Volume]] = None,
         gpu: typing.Union[None, str, modal.gpu._GPUConfig] = None,
         force_build: bool = False,
     ) -> Image:

modal/io_streams.py

@@ -21,7 +21,7 @@ from modal.exception import ClientClosed, ExecTimeoutError, InvalidError
 from modal_proto import api_pb2
 
 from ._utils.async_utils import synchronize_api
-from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 from ._utils.task_command_router_client import TaskCommandRouterClient
 from .client import _Client
 from .config import logger
@@ -64,7 +64,6 @@ async def _container_process_logs_iterator(
         get_raw_bytes=True,
         last_batch_index=last_index,
     )
-
     stream = client.stub.ContainerExecGetOutput.unary_stream(req)
     while True:
         # Check deadline before attempting to receive the next batch
@@ -76,11 +75,13 @@ async def _container_process_logs_iterator(
             break
         except StopAsyncIteration:
             break
+
+        for item in batch.items:
+            yield item.message_bytes, batch.batch_index
+
         if batch.HasField("exit_code"):
             yield None, batch.batch_index
             break
-        for item in batch.items:
-            yield item.message_bytes, batch.batch_index
 
 
 T = TypeVar("T", str, bytes)
@@ -89,7 +90,7 @@ T = TypeVar("T", str, bytes)
 class _StreamReaderThroughServer(Generic[T]):
     """A StreamReader implementation that reads from the server."""
 
-    _stream: Optional[AsyncGenerator[Optional[bytes], None]]
+    _stream: Optional[AsyncGenerator[T, None]]
 
     def __init__(
         self,
@@ -134,10 +135,9 @@ class _StreamReaderThroughServer(Generic[T]):
         self._stream_type = stream_type
 
         if self._object_type == "container_process":
-            # Container process streams need to be consumed as they are produced,
-            # otherwise the process will block. Use a buffer to store the stream
-            # until the client consumes it.
-            self._container_process_buffer: list[Optional[bytes]] = []
+            # TODO: we should not have this async code in constructors!
+            #  it only works as long as all the construction happens inside of synchronicity code
+            self._container_process_buffer: list[Optional[bytes]] = []  # TODO: change this to an asyncio.Queue
             self._consume_container_process_task = asyncio.create_task(self._consume_container_process_stream())
 
     @property
@@ -147,21 +147,18 @@ class _StreamReaderThroughServer(Generic[T]):
 
     async def read(self) -> T:
         """Fetch the entire contents of the stream until EOF."""
-        data_str = ""
-        data_bytes = b""
         logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read starting")
-        async for message in self._get_logs():
-            if message is None:
-                break
-            if self._text:
-                data_str += message.decode("utf-8")
-            else:
-                data_bytes += message
-
-        logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
         if self._text:
+            data_str = ""
+            async for message in _decode_bytes_stream_to_str(self._get_logs()):
+                data_str += message
+            logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
             return cast(T, data_str)
         else:
+            data_bytes = b""
+            async for message in self._get_logs():
+                data_bytes += message
+            logger.debug(f"{self._object_id} StreamReader fd={self._file_descriptor} read completed after EOF")
             return cast(T, data_bytes)
 
     async def _consume_container_process_stream(self):
@@ -181,6 +178,7 @@ class _StreamReaderThroughServer(Generic[T]):
                 )
                 async for message, batch_index in iterator:
                     if self._stream_type == StreamType.STDOUT and message:
+                        # TODO: rearchitect this, since these bytes aren't necessarily decodable
                         print(message.decode("utf-8"), end="")
                     elif self._stream_type == StreamType.PIPE:
                         self._container_process_buffer.append(message)
@@ -208,6 +206,9 @@ class _StreamReaderThroughServer(Generic[T]):
 
     async def _stream_container_process(self) -> AsyncGenerator[tuple[Optional[bytes], str], None]:
         """Streams the container process buffer to the reader."""
+        # Container process streams need to be consumed as they are produced,
+        # otherwise the process will block. Use a buffer to store the stream
+        # until the client consumes it.
         entry_id = 0
         if self._last_entry_id:
             entry_id = int(self._last_entry_id) + 1
@@ -225,7 +226,7 @@ class _StreamReaderThroughServer(Generic[T]):
 
             entry_id += 1
 
-    async def _get_logs(self, skip_empty_messages: bool = True) -> AsyncGenerator[Optional[bytes], None]:
+    async def _get_logs(self, skip_empty_messages: bool = True) -> AsyncGenerator[bytes, None]:
         """Streams sandbox or process logs from the server to the reader.
 
         Logs returned by this method may contain partial or multiple lines at a time.
@@ -237,7 +238,6 @@ class _StreamReaderThroughServer(Generic[T]):
             raise InvalidError("Logs can only be retrieved using the PIPE stream type.")
 
         if self.eof:
-            yield None
             return
 
         completed = False
@@ -262,6 +262,8 @@ class _StreamReaderThroughServer(Generic[T]):
                     if message is None:
                         completed = True
                         self.eof = True
+                        return
+
                     yield message
 
             except (GRPCError, StreamTerminatedError) as exc:
@@ -275,43 +277,37 @@ class _StreamReaderThroughServer(Generic[T]):
                         continue
                 raise
 
-    async def _get_logs_by_line(self) -> AsyncGenerator[Optional[bytes], None]:
+    async def _get_logs_by_line(self) -> AsyncGenerator[bytes, None]:
         """Process logs from the server and yield complete lines only."""
         async for message in self._get_logs():
-            if message is None:
-                if self._line_buffer:
-                    yield self._line_buffer
-                    self._line_buffer = b""
-                yield None
-            else:
-                assert isinstance(message, bytes)
-                self._line_buffer += message
-                while b"\n" in self._line_buffer:
-                    line, self._line_buffer = self._line_buffer.split(b"\n", 1)
-                    yield line + b"\n"
+            assert isinstance(message, bytes)
+            self._line_buffer += message
+            while b"\n" in self._line_buffer:
+                line, self._line_buffer = self._line_buffer.split(b"\n", 1)
+                yield line + b"\n"
+
+        if self._line_buffer:
+            yield self._line_buffer
+            self._line_buffer = b""
 
-    def _ensure_stream(self) -> AsyncGenerator[Optional[bytes], None]:
+    def _ensure_stream(self) -> AsyncGenerator[T, None]:
         if not self._stream:
             if self._by_line:
-                self._stream = self._get_logs_by_line()
+                # TODO: This is quite odd - it does line buffering in binary mode
+                # but we then always add the buffered text decoding on top of that.
+                # feels a bit upside down...
+                stream = self._get_logs_by_line()
             else:
-                self._stream = self._get_logs()
+                stream = self._get_logs()
+            if self._text:
+                stream = _decode_bytes_stream_to_str(stream)
+            self._stream = cast(AsyncGenerator[T, None], stream)
         return self._stream
 
     async def __anext__(self) -> T:
         """mdmd:hidden"""
         stream = self._ensure_stream()
-
-        value = await stream.__anext__()
-
-        # The stream yields None if it receives an EOF batch.
-        if value is None:
-            raise StopAsyncIteration
-
-        if self._text:
-            return cast(T, value.decode("utf-8"))
-        else:
-            return cast(T, value)
+        return cast(T, await stream.__anext__())
 
     async def aclose(self):
         """mdmd:hidden"""
@@ -330,6 +326,7 @@ async def _decode_bytes_stream_to_str(stream: AsyncGenerator[bytes, None]) -> As
         text = decoder.decode(item, final=False)
         if text:
             yield text
+
     # Flush any buffered partial character at end-of-stream
     tail = decoder.decode(b"", final=True)
     if tail:
@@ -495,23 +492,15 @@ class _StreamReader(Generic[T]):
 
     As an asynchronous iterable, the object supports the `for` and `async for`
     statements. Just loop over the object to read in chunks.
-
-    **Usage**
-
-    ```python fixture:running_app
-    from modal import Sandbox
-
-    sandbox = Sandbox.create(
-        "bash",
-        "-c",
-        "for i in $(seq 1 10); do echo foo; sleep 0.1; done",
-        app=running_app,
-    )
-    for message in sandbox.stdout:
-        print(f"Message: {message}")
-    ```
     """
 
+    _impl: Union[
+        _StreamReaderThroughServer,
+        _DevnullStreamReader,
+        _TextStreamReaderThroughCommandRouter,
+        _BytesStreamReaderThroughCommandRouter,
+    ]
+
     def __init__(
         self,
         file_descriptor: "api_pb2.FileDescriptor.ValueType",
@@ -570,19 +559,7 @@ class _StreamReader(Generic[T]):
         return self._impl.file_descriptor
 
     async def read(self) -> T:
-        """Fetch the entire contents of the stream until EOF.
-
-        **Usage**
-
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create("echo", "hello", app=running_app)
-        sandbox.wait()
-
-        print(sandbox.stdout.read())
-        ```
-        """
+        """Fetch the entire contents of the stream until EOF."""
         return await self._impl.read()
 
     # TODO(saltzm): I'd prefer to have the implementation classes only implement __aiter__
@@ -660,15 +637,13 @@ class _StreamWriterThroughServer:
 
         try:
             if self._object_type == "sandbox":
-                await retry_transient_errors(
-                    self._client.stub.SandboxStdinWrite,
+                await self._client.stub.SandboxStdinWrite(
                     api_pb2.SandboxStdinWriteRequest(
                         sandbox_id=self._object_id, index=index, eof=self._is_closed, input=data
                     ),
                 )
             else:
-                await retry_transient_errors(
-                    self._client.stub.ContainerExecPutInput,
+                await self._client.stub.ContainerExecPutInput(
                     api_pb2.ContainerExecPutInputRequest(
                         exec_id=self._object_id,
                         input=api_pb2.RuntimeInputMessage(message=data, message_index=index, eof=self._is_closed),
@@ -753,21 +728,16 @@ class _StreamWriter:
 
         **Usage**
 
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create(
+        ```python fixture:sandbox
+        proc = sandbox.exec(
             "bash",
             "-c",
             "while read line; do echo $line; done",
-            app=running_app,
         )
-        sandbox.stdin.write(b"foo\\n")
-        sandbox.stdin.write(b"bar\\n")
-        sandbox.stdin.write_eof()
-
-        sandbox.stdin.drain()
-        sandbox.wait()
+        proc.stdin.write(b"foo\\n")
+        proc.stdin.write(b"bar\\n")
+        proc.stdin.write_eof()
+        proc.stdin.drain()
         ```
         """
         self._impl.write(data)

modal/io_streams.pyi

@@ -21,7 +21,7 @@ T = typing.TypeVar("T")
 class _StreamReaderThroughServer(typing.Generic[T]):
     """A StreamReader implementation that reads from the server."""
 
-    _stream: typing.Optional[collections.abc.AsyncGenerator[typing.Optional[bytes], None]]
+    _stream: typing.Optional[collections.abc.AsyncGenerator[T, None]]
 
     def __init__(
         self,
@@ -54,9 +54,7 @@ class _StreamReaderThroughServer(typing.Generic[T]):
         """Streams the container process buffer to the reader."""
         ...
 
-    def _get_logs(
-        self, skip_empty_messages: bool = True
-    ) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]:
+    def _get_logs(self, skip_empty_messages: bool = True) -> collections.abc.AsyncGenerator[bytes, None]:
         """Streams sandbox or process logs from the server to the reader.
 
         Logs returned by this method may contain partial or multiple lines at a time.
@@ -66,11 +64,11 @@ class _StreamReaderThroughServer(typing.Generic[T]):
         """
         ...
 
-    def _get_logs_by_line(self) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]:
+    def _get_logs_by_line(self) -> collections.abc.AsyncGenerator[bytes, None]:
         """Process logs from the server and yield complete lines only."""
         ...
 
-    def _ensure_stream(self) -> collections.abc.AsyncGenerator[typing.Optional[bytes], None]: ...
+    def _ensure_stream(self) -> collections.abc.AsyncGenerator[T, None]: ...
     async def __anext__(self) -> T:
         """mdmd:hidden"""
         ...
@@ -184,22 +182,15 @@ class _StreamReader(typing.Generic[T]):
 
     As an asynchronous iterable, the object supports the `for` and `async for`
     statements. Just loop over the object to read in chunks.
+    """
 
-    **Usage**
-
-    ```python fixture:running_app
-    from modal import Sandbox
+    _impl: typing.Union[
+        _StreamReaderThroughServer,
+        _DevnullStreamReader,
+        _TextStreamReaderThroughCommandRouter,
+        _BytesStreamReaderThroughCommandRouter,
+    ]
 
-    sandbox = Sandbox.create(
-        "bash",
-        "-c",
-        "for i in $(seq 1 10); do echo foo; sleep 0.1; done",
-        app=running_app,
-    )
-    for message in sandbox.stdout:
-        print(f"Message: {message}")
-    ```
-    """
     def __init__(
         self,
         file_descriptor: int,
@@ -222,19 +213,7 @@ class _StreamReader(typing.Generic[T]):
         ...
 
     async def read(self) -> T:
-        """Fetch the entire contents of the stream until EOF.
-
-        **Usage**
-
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create("echo", "hello", app=running_app)
-        sandbox.wait()
-
-        print(sandbox.stdout.read())
-        ```
-        """
+        """Fetch the entire contents of the stream until EOF."""
         ...
 
     def __aiter__(self) -> collections.abc.AsyncIterator[T]:
@@ -318,21 +297,16 @@ class _StreamWriter:
 
         **Usage**
 
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create(
+        ```python fixture:sandbox
+        proc = sandbox.exec(
             "bash",
             "-c",
             "while read line; do echo $line; done",
-            app=running_app,
         )
-        sandbox.stdin.write(b"foo\n")
-        sandbox.stdin.write(b"bar\n")
-        sandbox.stdin.write_eof()
-
-        sandbox.stdin.drain()
-        sandbox.wait()
+        proc.stdin.write(b"foo\n")
+        proc.stdin.write(b"bar\n")
+        proc.stdin.write_eof()
+        proc.stdin.drain()
         ```
         """
         ...
@@ -376,22 +350,15 @@ class StreamReader(typing.Generic[T]):
 
     As an asynchronous iterable, the object supports the `for` and `async for`
     statements. Just loop over the object to read in chunks.
+    """
 
-    **Usage**
-
-    ```python fixture:running_app
-    from modal import Sandbox
+    _impl: typing.Union[
+        _StreamReaderThroughServer,
+        _DevnullStreamReader,
+        _TextStreamReaderThroughCommandRouter,
+        _BytesStreamReaderThroughCommandRouter,
+    ]
 
-    sandbox = Sandbox.create(
-        "bash",
-        "-c",
-        "for i in $(seq 1 10); do echo foo; sleep 0.1; done",
-        app=running_app,
-    )
-    for message in sandbox.stdout:
-        print(f"Message: {message}")
-    ```
-    """
     def __init__(
         self,
         file_descriptor: int,
@@ -415,35 +382,11 @@ class StreamReader(typing.Generic[T]):
 
     class __read_spec(typing_extensions.Protocol[T_INNER, SUPERSELF]):
         def __call__(self, /) -> T_INNER:
-            """Fetch the entire contents of the stream until EOF.
-
-            **Usage**
-
-            ```python fixture:running_app
-            from modal import Sandbox
-
-            sandbox = Sandbox.create("echo", "hello", app=running_app)
-            sandbox.wait()
-
-            print(sandbox.stdout.read())
-            ```
-            """
+            """Fetch the entire contents of the stream until EOF."""
             ...
 
         async def aio(self, /) -> T_INNER:
-            """Fetch the entire contents of the stream until EOF.
-
-            **Usage**
-
-            ```python fixture:running_app
-            from modal import Sandbox
-
-            sandbox = Sandbox.create("echo", "hello", app=running_app)
-            sandbox.wait()
-
-            print(sandbox.stdout.read())
-            ```
-            """
+            """Fetch the entire contents of the stream until EOF."""
             ...
 
     read: __read_spec[T, typing_extensions.Self]
@@ -493,21 +436,16 @@ class StreamWriter:
 
         **Usage**
 
-        ```python fixture:running_app
-        from modal import Sandbox
-
-        sandbox = Sandbox.create(
+        ```python fixture:sandbox
+        proc = sandbox.exec(
             "bash",
             "-c",
             "while read line; do echo $line; done",
-            app=running_app,
         )
-        sandbox.stdin.write(b"foo\n")
-        sandbox.stdin.write(b"bar\n")
-        sandbox.stdin.write_eof()
-
-        sandbox.stdin.drain()
-        sandbox.wait()
+        proc.stdin.write(b"foo\n")
+        proc.stdin.write(b"bar\n")
+        proc.stdin.write_eof()
+        proc.stdin.drain()
         ```
         """
         ...

modal/mount.py

@@ -24,7 +24,7 @@ from ._object import _get_environment_name, _Object
 from ._resolver import Resolver
 from ._utils.async_utils import TaskContext, aclosing, async_map, synchronize_api
 from ._utils.blob_utils import FileUploadSpec, blob_upload_file, get_file_upload_spec_from_path
-from ._utils.grpc_utils import retry_transient_errors
+from ._utils.grpc_utils import Retry
 from ._utils.name_utils import check_object_name
 from ._utils.package_utils import get_module_mount_info
 from .client import _Client
@@ -518,7 +518,7 @@ class _Mount(_Object, type_prefix="mo"):
 
             request = api_pb2.MountPutFileRequest(sha256_hex=file_spec.sha256_hex)
             accounted_hashes.add(file_spec.sha256_hex)
-            response = await retry_transient_errors(resolver.client.stub.MountPutFile, request, base_delay=1)
+            response = await resolver.client.stub.MountPutFile(request, retry=Retry(base_delay=1))
 
             if response.exists:
                 n_finished += 1
@@ -544,7 +544,7 @@ class _Mount(_Object, type_prefix="mo"):
 
             start_time = time.monotonic()
             while time.monotonic() - start_time < MOUNT_PUT_FILE_CLIENT_TIMEOUT:
-                response = await retry_transient_errors(resolver.client.stub.MountPutFile, request2, base_delay=1)
+                response = await resolver.client.stub.MountPutFile(request2, retry=Retry(base_delay=1))
                 if response.exists:
                     n_finished += 1
                     return mount_file
@@ -591,7 +591,7 @@ class _Mount(_Object, type_prefix="mo"):
                 environment_name=resolver.environment_name,
             )
 
-        resp = await retry_transient_errors(resolver.client.stub.MountGetOrCreate, req, base_delay=1)
+        resp = await resolver.client.stub.MountGetOrCreate(req, retry=Retry(base_delay=1))
         status_row.finish(f"Created mount {message_label}")
 
         logger.debug(f"Uploaded {total_uploads} new files and {total_bytes} bytes in {time.monotonic() - t0}s")