modal 0.67.42__py3-none-any.whl → 0.68.4__py3-none-any.whl

modal/file_io.py

@@ -0,0 +1,380 @@
+# Copyright Modal Labs 2024
+import asyncio
+import io
+from typing import TYPE_CHECKING, AsyncIterator, Generic, Optional, Sequence, TypeVar, Union, cast
+
+if TYPE_CHECKING:
+    import _typeshed
+
+from grpclib.exceptions import GRPCError, StreamTerminatedError
+
+from modal._utils.grpc_utils import retry_transient_errors
+from modal_proto import api_pb2
+
+from ._utils.async_utils import synchronize_api
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES
+from .client import _Client
+from .exception import FilesystemExecutionError, InvalidError
+
+LARGE_FILE_SIZE_LIMIT = 16 * 1024 * 1024  # 16 MiB
+READ_FILE_SIZE_LIMIT = 100 * 1024 * 1024  # 100 MiB
+
+ERROR_MAPPING = {
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_UNSPECIFIED: FilesystemExecutionError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_PERM: PermissionError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_NOENT: FileNotFoundError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_IO: IOError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_NXIO: IOError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_NOMEM: MemoryError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_ACCES: PermissionError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_EXIST: FileExistsError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_NOTDIR: NotADirectoryError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_ISDIR: IsADirectoryError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_INVAL: OSError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_MFILE: OSError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_FBIG: OSError,
+    api_pb2.SystemErrorCode.SYSTEM_ERROR_CODE_NOSPC: OSError,
+}
+
+T = TypeVar("T", str, bytes)
+
+
+async def _delete_bytes(file: "_FileIO", start: Optional[int] = None, end: Optional[int] = None) -> None:
+    """Delete a range of bytes from the file.
+
+    `start` and `end` are byte offsets. `start` is inclusive, `end` is exclusive.
+    If either is None, the start or end of the file is used, respectively.
+    """
+    assert file._file_descriptor is not None
+    file._check_closed()
+    if start is not None and end is not None:
+        if start >= end:
+            raise ValueError("start must be less than end")
+    resp = await file._make_request(
+        api_pb2.ContainerFilesystemExecRequest(
+            file_delete_bytes_request=api_pb2.ContainerFileDeleteBytesRequest(
+                file_descriptor=file._file_descriptor,
+                start_inclusive=start,
+                end_exclusive=end,
+            ),
+            task_id=file._task_id,
+        )
+    )
+    await file._wait(resp.exec_id)
+
+
+async def _replace_bytes(file: "_FileIO", data: bytes, start: Optional[int] = None, end: Optional[int] = None) -> None:
+    """Replace a range of bytes in the file with new data. The length of the data does not
+    have to be the same as the length of the range being replaced.
+
+    `start` and `end` are byte offsets. `start` is inclusive, `end` is exclusive.
+    If either is None, the start or end of the file is used, respectively.
+    """
+    assert file._file_descriptor is not None
+    file._check_closed()
+    if start is not None and end is not None:
+        if start >= end:
+            raise InvalidError("start must be less than end")
+    if len(data) > LARGE_FILE_SIZE_LIMIT:
+        raise InvalidError("Write request payload exceeds 16 MiB limit")
+    resp = await file._make_request(
+        api_pb2.ContainerFilesystemExecRequest(
+            file_write_replace_bytes_request=api_pb2.ContainerFileWriteReplaceBytesRequest(
+                file_descriptor=file._file_descriptor,
+                data=data,
+                start_inclusive=start,
+                end_exclusive=end,
+            ),
+            task_id=file._task_id,
+        )
+    )
+    await file._wait(resp.exec_id)
+
+
+# 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]):
+    """FileIO handle, used in the Sandbox filesystem API.
+
+    The API is designed to mimic Python's io.FileIO.
+
+    **Usage**
+
+    ```python
+    import modal
+
+    app = modal.App.lookup("my-app", create_if_missing=True)
+
+    sb = modal.Sandbox.create(app=app)
+    f = sb.open("/tmp/foo.txt", "w")
+    f.write("hello")
+    f.close()
+    ```
+    """
+
+    _binary = False
+    _readable = False
+    _writable = False
+    _appended = False
+    _closed = True
+
+    _task_id: str = ""
+    _file_descriptor: str = ""
+    _client: Optional[_Client] = None
+
+    def _validate_mode(self, mode: str) -> None:
+        if not any(char in mode for char in "rwax"):
+            raise ValueError(f"Invalid file mode: {mode}")
+
+        self._readable = "r" in mode or "+" in mode
+        self._writable = "w" in mode or "a" in mode or "x" in mode or "+" in mode
+        self._appended = "a" in mode
+        self._binary = "b" in mode
+
+        valid_chars = set("rwaxb+")
+        if any(char not in valid_chars for char in mode):
+            raise ValueError(f"Invalid file mode: {mode}")
+
+        mode_count = sum(1 for c in mode if c in "rwax")
+        if mode_count > 1:
+            raise ValueError("must have exactly one of create/read/write/append mode")
+
+        seen_chars = set()
+        for char in mode:
+            if char in seen_chars:
+                raise ValueError(f"Invalid file mode: {mode}")
+            seen_chars.add(char)
+
+    def _handle_error(self, error: api_pb2.SystemErrorMessage) -> None:
+        error_class = ERROR_MAPPING.get(error.error_code, FilesystemExecutionError)
+        raise error_class(error.error_message)
+
+    async def _consume_output(self, exec_id: str) -> AsyncIterator[Optional[bytes]]:
+        req = api_pb2.ContainerFilesystemExecGetOutputRequest(
+            exec_id=exec_id,
+            timeout=55,
+        )
+        assert self._client is not None
+        async for batch in self._client.stub.ContainerFilesystemExecGetOutput.unary_stream(req):
+            if batch.eof:
+                yield None
+                break
+            if batch.HasField("error"):
+                self._handle_error(batch.error)
+            for message in batch.output:
+                yield message
+
+    async def _wait(self, exec_id: str) -> bytes:
+        # The logic here is similar to how output is read from `exec`
+        output = b""
+        completed = False
+        retries_remaining = 10
+        while not completed:
+            try:
+                async for data in self._consume_output(exec_id):
+                    if data is None:
+                        completed = True
+                        break
+                    output += data
+            except (GRPCError, StreamTerminatedError) 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
+                raise
+        return output
+
+    def _validate_type(self, data: Union[bytes, str]) -> None:
+        if self._binary and isinstance(data, str):
+            raise TypeError("Expected bytes when in binary mode")
+        if not self._binary and isinstance(data, bytes):
+            raise TypeError("Expected str when in text mode")
+
+    async def _open_file(self, path: str, mode: str) -> None:
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_open_request=api_pb2.ContainerFileOpenRequest(path=path, mode=mode),
+                task_id=self._task_id,
+            )
+        )
+        if not resp.HasField("file_descriptor"):
+            raise FilesystemExecutionError("Failed to open file")
+        self._file_descriptor = resp.file_descriptor
+        await self._wait(resp.exec_id)
+
+    @classmethod
+    async def create(
+        cls, path: str, mode: Union["_typeshed.OpenTextMode", "_typeshed.OpenBinaryMode"], client: _Client, task_id: str
+    ) -> "_FileIO":
+        """Create a new FileIO handle."""
+        self = cls.__new__(cls)
+        self._client = client
+        self._task_id = task_id
+        self._validate_mode(mode)
+        await self._open_file(path, mode)
+        self._closed = False
+        return self
+
+    async def _make_request(
+        self, request: api_pb2.ContainerFilesystemExecRequest
+    ) -> api_pb2.ContainerFilesystemExecResponse:
+        assert self._client is not None
+        return await retry_transient_errors(self._client.stub.ContainerFilesystemExec, request)
+
+    async def _make_read_request(self, n: Optional[int]) -> bytes:
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_read_request=api_pb2.ContainerFileReadRequest(file_descriptor=self._file_descriptor, n=n),
+                task_id=self._task_id,
+            )
+        )
+        return await self._wait(resp.exec_id)
+
+    async def read(self, n: Optional[int] = None) -> T:
+        """Read n bytes from the current position, or the entire remaining file if n is None."""
+        self._check_closed()
+        self._check_readable()
+        if n is not None and n > READ_FILE_SIZE_LIMIT:
+            raise ValueError("Read request payload exceeds 100 MiB limit")
+        output = await self._make_read_request(n)
+        if self._binary:
+            return cast(T, output)
+        return cast(T, output.decode("utf-8"))
+
+    async def readline(self) -> T:
+        """Read a single line from the current position."""
+        self._check_closed()
+        self._check_readable()
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_read_line_request=api_pb2.ContainerFileReadLineRequest(file_descriptor=self._file_descriptor),
+                task_id=self._task_id,
+            )
+        )
+        output = await self._wait(resp.exec_id)
+        if self._binary:
+            return cast(T, output)
+        return cast(T, output.decode("utf-8"))
+
+    async def readlines(self) -> Sequence[T]:
+        """Read all lines from the current position."""
+        self._check_closed()
+        self._check_readable()
+        output = await self._make_read_request(None)
+        if self._binary:
+            lines_bytes = output.split(b"\n")
+            output = [line + b"\n" for line in lines_bytes[:-1]] + ([lines_bytes[-1]] if lines_bytes[-1] else [])
+            return cast(Sequence[T], output)
+        else:
+            lines = output.decode("utf-8").split("\n")
+            output = [line + "\n" for line in lines[:-1]] + ([lines[-1]] if lines[-1] else [])
+            return cast(Sequence[T], output)
+
+    async def write(self, data: Union[bytes, str]) -> None:
+        """Write data to the current position.
+
+        Writes may not appear until the entire buffer is flushed, which
+        can be done manually with `flush()` or automatically when the file is
+        closed.
+        """
+        self._check_closed()
+        self._check_writable()
+        self._validate_type(data)
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        if len(data) > LARGE_FILE_SIZE_LIMIT:
+            raise ValueError("Write request payload exceeds 16 MiB limit")
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_write_request=api_pb2.ContainerFileWriteRequest(file_descriptor=self._file_descriptor, data=data),
+                task_id=self._task_id,
+            )
+        )
+        await self._wait(resp.exec_id)
+
+    async def flush(self) -> None:
+        """Flush the buffer to disk."""
+        self._check_closed()
+        self._check_writable()
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_flush_request=api_pb2.ContainerFileFlushRequest(file_descriptor=self._file_descriptor),
+                task_id=self._task_id,
+            )
+        )
+        await self._wait(resp.exec_id)
+
+    def _get_whence(self, whence: int):
+        if whence == 0:
+            return api_pb2.SeekWhence.SEEK_SET
+        elif whence == 1:
+            return api_pb2.SeekWhence.SEEK_CUR
+        elif whence == 2:
+            return api_pb2.SeekWhence.SEEK_END
+        else:
+            raise ValueError(f"Invalid whence value: {whence}")
+
+    async def seek(self, offset: int, whence: int = 0) -> None:
+        """Move to a new position in the file.
+
+        `whence` defaults to 0 (absolute file positioning); other values are 1
+        (relative to the current position) and 2 (relative to the file's end).
+        """
+        self._check_closed()
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_seek_request=api_pb2.ContainerFileSeekRequest(
+                    file_descriptor=self._file_descriptor,
+                    offset=offset,
+                    whence=self._get_whence(whence),
+                ),
+                task_id=self._task_id,
+            )
+        )
+        await self._wait(resp.exec_id)
+
+    async def _close(self) -> None:
+        # Buffer is flushed by the runner on close
+        resp = await self._make_request(
+            api_pb2.ContainerFilesystemExecRequest(
+                file_close_request=api_pb2.ContainerFileCloseRequest(file_descriptor=self._file_descriptor),
+                task_id=self._task_id,
+            )
+        )
+        self._closed = True
+        await self._wait(resp.exec_id)
+
+    async def close(self) -> None:
+        """Flush the buffer and close the file."""
+        await self._close()
+
+    # also validated in the runner, but checked in the client to catch errors early
+    def _check_writable(self) -> None:
+        if not self._writable:
+            raise io.UnsupportedOperation("not writeable")
+
+    # also validated in the runner, but checked in the client to catch errors early
+    def _check_readable(self) -> None:
+        if not self._readable:
+            raise io.UnsupportedOperation("not readable")
+
+    # also validated in the runner, but checked in the client to catch errors early
+    def _check_closed(self) -> None:
+        if self._closed:
+            raise ValueError("I/O operation on closed file")
+
+    def __enter__(self) -> "_FileIO":
+        self._check_closed()
+        return self
+
+    async def __exit__(self, exc_type, exc_value, traceback) -> None:
+        await self._close()
+
+
+delete_bytes = synchronize_api(_delete_bytes)
+replace_bytes = synchronize_api(_replace_bytes)
+FileIO = synchronize_api(_FileIO)

modal/file_io.pyi

@@ -0,0 +1,185 @@
+import _typeshed
+import modal.client
+import modal_proto.api_pb2
+import typing
+import typing_extensions
+
+T = typing.TypeVar("T")
+
+async def _delete_bytes(
+    file: _FileIO, start: typing.Optional[int] = None, end: typing.Optional[int] = None
+) -> None: ...
+async def _replace_bytes(
+    file: _FileIO, data: bytes, start: typing.Optional[int] = None, end: typing.Optional[int] = None
+) -> None: ...
+
+class _FileIO(typing.Generic[T]):
+    _task_id: str
+    _file_descriptor: str
+    _client: typing.Optional[modal.client._Client]
+
+    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 _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: ...
+    @classmethod
+    async def create(
+        cls,
+        path: str,
+        mode: typing.Union[_typeshed.OpenTextMode, _typeshed.OpenBinaryMode],
+        client: modal.client._Client,
+        task_id: str,
+    ) -> _FileIO: ...
+    async def _make_request(
+        self, request: modal_proto.api_pb2.ContainerFilesystemExecRequest
+    ) -> modal_proto.api_pb2.ContainerFilesystemExecResponse: ...
+    async def _make_read_request(self, n: typing.Optional[int]) -> bytes: ...
+    async def read(self, n: typing.Optional[int] = None) -> T: ...
+    async def readline(self) -> T: ...
+    async def readlines(self) -> typing.Sequence[T]: ...
+    async def write(self, data: typing.Union[bytes, str]) -> None: ...
+    async def flush(self) -> None: ...
+    def _get_whence(self, whence: int): ...
+    async def seek(self, offset: int, whence: int = 0) -> None: ...
+    async def _close(self) -> None: ...
+    async def close(self) -> None: ...
+    def _check_writable(self) -> None: ...
+    def _check_readable(self) -> None: ...
+    def _check_closed(self) -> None: ...
+    def __enter__(self) -> _FileIO: ...
+    async def __exit__(self, exc_type, exc_value, traceback) -> None: ...
+
+class __delete_bytes_spec(typing_extensions.Protocol):
+    def __call__(self, file: FileIO, start: typing.Optional[int] = None, end: typing.Optional[int] = None) -> None: ...
+    async def aio(self, file: FileIO, start: typing.Optional[int] = None, end: typing.Optional[int] = None) -> None: ...
+
+delete_bytes: __delete_bytes_spec
+
+class __replace_bytes_spec(typing_extensions.Protocol):
+    def __call__(
+        self, file: FileIO, data: bytes, start: typing.Optional[int] = None, end: typing.Optional[int] = None
+    ) -> None: ...
+    async def aio(
+        self, file: FileIO, data: bytes, start: typing.Optional[int] = None, end: typing.Optional[int] = None
+    ) -> None: ...
+
+replace_bytes: __replace_bytes_spec
+
+T_INNER = typing.TypeVar("T_INNER", covariant=True)
+
+class FileIO(typing.Generic[T]):
+    _task_id: str
+    _file_descriptor: str
+    _client: typing.Optional[modal.client.Client]
+
+    def __init__(self, /, *args, **kwargs): ...
+    def _validate_mode(self, mode: str) -> None: ...
+    def _handle_error(self, error: modal_proto.api_pb2.SystemErrorMessage) -> None: ...
+
+    class ___consume_output_spec(typing_extensions.Protocol):
+        def __call__(self, exec_id: str) -> typing.Iterator[typing.Optional[bytes]]: ...
+        def aio(self, exec_id: str) -> typing.AsyncIterator[typing.Optional[bytes]]: ...
+
+    _consume_output: ___consume_output_spec
+
+    class ___wait_spec(typing_extensions.Protocol):
+        def __call__(self, exec_id: str) -> bytes: ...
+        async def aio(self, exec_id: str) -> bytes: ...
+
+    _wait: ___wait_spec
+
+    def _validate_type(self, data: typing.Union[bytes, str]) -> None: ...
+
+    class ___open_file_spec(typing_extensions.Protocol):
+        def __call__(self, path: str, mode: str) -> None: ...
+        async def aio(self, path: str, mode: str) -> None: ...
+
+    _open_file: ___open_file_spec
+
+    @classmethod
+    def create(
+        cls,
+        path: str,
+        mode: typing.Union[_typeshed.OpenTextMode, _typeshed.OpenBinaryMode],
+        client: modal.client.Client,
+        task_id: str,
+    ) -> FileIO: ...
+
+    class ___make_request_spec(typing_extensions.Protocol):
+        def __call__(
+            self, request: modal_proto.api_pb2.ContainerFilesystemExecRequest
+        ) -> modal_proto.api_pb2.ContainerFilesystemExecResponse: ...
+        async def aio(
+            self, request: modal_proto.api_pb2.ContainerFilesystemExecRequest
+        ) -> modal_proto.api_pb2.ContainerFilesystemExecResponse: ...
+
+    _make_request: ___make_request_spec
+
+    class ___make_read_request_spec(typing_extensions.Protocol):
+        def __call__(self, n: typing.Optional[int]) -> bytes: ...
+        async def aio(self, n: typing.Optional[int]) -> bytes: ...
+
+    _make_read_request: ___make_read_request_spec
+
+    class __read_spec(typing_extensions.Protocol[T_INNER]):
+        def __call__(self, n: typing.Optional[int] = None) -> T_INNER: ...
+        async def aio(self, n: typing.Optional[int] = None) -> T_INNER: ...
+
+    read: __read_spec[T]
+
+    class __readline_spec(typing_extensions.Protocol[T_INNER]):
+        def __call__(self) -> T_INNER: ...
+        async def aio(self) -> T_INNER: ...
+
+    readline: __readline_spec[T]
+
+    class __readlines_spec(typing_extensions.Protocol[T_INNER]):
+        def __call__(self) -> typing.Sequence[T_INNER]: ...
+        async def aio(self) -> typing.Sequence[T_INNER]: ...
+
+    readlines: __readlines_spec[T]
+
+    class __write_spec(typing_extensions.Protocol):
+        def __call__(self, data: typing.Union[bytes, str]) -> None: ...
+        async def aio(self, data: typing.Union[bytes, str]) -> None: ...
+
+    write: __write_spec
+
+    class __flush_spec(typing_extensions.Protocol):
+        def __call__(self) -> None: ...
+        async def aio(self) -> None: ...
+
+    flush: __flush_spec
+
+    def _get_whence(self, whence: int): ...
+
+    class __seek_spec(typing_extensions.Protocol):
+        def __call__(self, offset: int, whence: int = 0) -> None: ...
+        async def aio(self, offset: int, whence: int = 0) -> None: ...
+
+    seek: __seek_spec
+
+    class ___close_spec(typing_extensions.Protocol):
+        def __call__(self) -> None: ...
+        async def aio(self) -> None: ...
+
+    _close: ___close_spec
+
+    class __close_spec(typing_extensions.Protocol):
+        def __call__(self) -> None: ...
+        async def aio(self) -> None: ...
+
+    close: __close_spec
+
+    def _check_writable(self) -> None: ...
+    def _check_readable(self) -> None: ...
+    def _check_closed(self) -> None: ...
+    def __enter__(self) -> FileIO: ...
+
+    class ____exit___spec(typing_extensions.Protocol):
+        def __call__(self, exc_type, exc_value, traceback) -> None: ...
+        async def aio(self, exc_type, exc_value, traceback) -> None: ...
+
+    __exit__: ____exit___spec

modal/functions.py

@@ -32,6 +32,7 @@ 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 ._traceback import print_server_warnings
 from ._utils.async_utils import (
     TaskContext,
     async_merge,
@@ -1061,6 +1062,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 else:
                     raise
 
+            print_server_warnings(response.server_warnings)
+
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
         rep = f"Ref({app_name})"

modal/functions.pyi

@@ -455,11 +455,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]: ...
@@ -471,17 +471,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/partial_function.py

@@ -49,7 +49,7 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
     raw_f: Callable[P, ReturnType]
     flags: _PartialFunctionFlags
     webhook_config: Optional[api_pb2.WebhookConfig]
-    is_generator: Optional[bool]
+    is_generator: bool
     keep_warm: Optional[int]
     batch_max_size: Optional[int]
     batch_wait_ms: Optional[int]
@@ -73,7 +73,13 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
         self.raw_f = raw_f
         self.flags = flags
         self.webhook_config = webhook_config
-        self.is_generator = is_generator
+        if is_generator is None:
+            # auto detect - doesn't work if the function *returns* a generator
+            final_is_generator = inspect.isgeneratorfunction(raw_f) or inspect.isasyncgenfunction(raw_f)
+        else:
+            final_is_generator = is_generator
+
+        self.is_generator = final_is_generator
         self.keep_warm = keep_warm
         self.wrapped = False  # Make sure that this was converted into a FunctionHandle
         self.batch_max_size = batch_max_size
@@ -101,7 +107,7 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
             # This happens mainly during serialization of the wrapped underlying class of a Cls
             # since we don't have the instance info here we just return the PartialFunction itself
             # to let it be bound to a variable and become a Function later on
-            return self
+            return self  # type: ignore  # this returns a PartialFunction in a special internal case
 
     def __del__(self):
         if (self.flags & _PartialFunctionFlags.FUNCTION) and self.wrapped is False:
@@ -154,14 +160,14 @@ def _find_partial_methods_for_user_cls(user_cls: type[Any], flags: int) -> dict[
             )
             deprecation_error((2024, 2, 21), message)
 
-    partial_functions: dict[str, PartialFunction] = {}
+    partial_functions: dict[str, _PartialFunction] = {}
     for parent_cls in reversed(user_cls.mro()):
         if parent_cls is not object:
             for k, v in parent_cls.__dict__.items():
-                if isinstance(v, PartialFunction):
-                    partial_function = synchronizer._translate_in(v)  # TODO: remove need for?
-                    if partial_function.flags & flags:
-                        partial_functions[k] = partial_function
+                if isinstance(v, PartialFunction):  # type: ignore[reportArgumentType]   # synchronicity wrapper types
+                    _partial_function: _PartialFunction = typing.cast(_PartialFunction, synchronizer._translate_in(v))
+                    if _partial_function.flags & flags:
+                        partial_functions[k] = _partial_function
 
     return partial_functions
 
@@ -245,8 +251,6 @@ def _method(
                 "Batched function on classes should not be wrapped by `@method`. "
                 "Suggestion: remove the `@method` decorator."
             )
-        if is_generator is None:
-            is_generator = inspect.isgeneratorfunction(raw_f) or inspect.isasyncgenfunction(raw_f)
         return _PartialFunction(raw_f, _PartialFunctionFlags.FUNCTION, is_generator=is_generator, keep_warm=keep_warm)
 
     return wrapper