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

modal/file_io.py

@@ -0,0 +1,537 @@
+# 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:
+    import _typeshed
+
+import json
+
+from grpclib.exceptions import GRPCError, StreamTerminatedError
+
+from modal._utils.async_utils import TaskContext
+from modal._utils.grpc_utils import retry_transient_errors
+from modal.exception import ClientClosed
+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
+
+WRITE_CHUNK_SIZE = 16 * 1024 * 1024  # 16 MiB
+WRITE_FILE_SIZE_LIMIT = 1024 * 1024 * 1024  # 1 GiB
+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 retry_transient_errors(
+        file._client.stub.ContainerFilesystemExec,
+        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) > WRITE_CHUNK_SIZE:
+        raise InvalidError("Write request payload exceeds 16 MiB limit")
+    resp = await retry_transient_errors(
+        file._client.stub.ContainerFilesystemExec,
+        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)
+
+
+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]):
+    """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: _Client
+    _watch_output_buffer: list[Optional[bytes]] = []
+
+    def __init__(self, client: _Client, task_id: str) -> None:
+        self._client = client
+        self._task_id = task_id
+
+    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,
+        )
+        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 _consume_watch_output(self, exec_id: str) -> None:
+        completed = False
+        retries_remaining = 10
+        while not completed:
+            try:
+                iterator = self._consume_output(exec_id)
+                async for message in iterator:
+                    self._watch_output_buffer.append(message)
+                    if message is None:
+                        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(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
+
+    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 _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 retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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 = _FileIO(client, task_id)
+        self._validate_mode(mode)
+        await self._open_file(path, mode)
+        self._closed = False
+        return self
+
+    async def _make_read_request(self, n: Optional[int]) -> bytes:
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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 retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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")
+            return_bytes = [line + b"\n" for line in lines_bytes[:-1]] + ([lines_bytes[-1]] if lines_bytes[-1] else [])
+            return cast(Sequence[T], return_bytes)
+        else:
+            lines = output.decode("utf-8").split("\n")
+            return_strs = [line + "\n" for line in lines[:-1]] + ([lines[-1]] if lines[-1] else [])
+            return cast(Sequence[T], return_strs)
+
+    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) > WRITE_FILE_SIZE_LIMIT:
+            raise ValueError("Write request payload exceeds 1 GiB limit")
+        for i in range(0, len(data), WRITE_CHUNK_SIZE):
+            chunk = data[i : i + WRITE_CHUNK_SIZE]
+            resp = await retry_transient_errors(
+                self._client.stub.ContainerFilesystemExec,
+                api_pb2.ContainerFilesystemExecRequest(
+                    file_write_request=api_pb2.ContainerFileWriteRequest(
+                        file_descriptor=self._file_descriptor,
+                        data=chunk,
+                    ),
+                    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 retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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 retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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)
+
+    @classmethod
+    async def ls(cls, path: str, client: _Client, task_id: str) -> list[str]:
+        """List the contents of the provided directory."""
+        self = _FileIO(client, task_id)
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            api_pb2.ContainerFilesystemExecRequest(
+                file_ls_request=api_pb2.ContainerFileLsRequest(path=path),
+                task_id=task_id,
+            ),
+        )
+        output = await self._wait(resp.exec_id)
+        try:
+            return json.loads(output.decode("utf-8"))["paths"]
+        except json.JSONDecodeError:
+            raise FilesystemExecutionError("failed to parse list output")
+
+    @classmethod
+    async def mkdir(cls, path: str, client: _Client, task_id: str, parents: bool = False) -> None:
+        """Create a new directory."""
+        self = _FileIO(client, task_id)
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            api_pb2.ContainerFilesystemExecRequest(
+                file_mkdir_request=api_pb2.ContainerFileMkdirRequest(path=path, make_parents=parents),
+                task_id=self._task_id,
+            ),
+        )
+        await self._wait(resp.exec_id)
+
+    @classmethod
+    async def rm(cls, path: str, client: _Client, task_id: str, recursive: bool = False) -> None:
+        """Remove a file or directory in the Sandbox."""
+        self = _FileIO(client, task_id)
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            api_pb2.ContainerFilesystemExecRequest(
+                file_rm_request=api_pb2.ContainerFileRmRequest(path=path, recursive=recursive),
+                task_id=self._task_id,
+            ),
+        )
+        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 = _FileIO(client, task_id)
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            api_pb2.ContainerFilesystemExecRequest(
+                file_watch_request=api_pb2.ContainerFileWatchRequest(
+                    path=path,
+                    recursive=recursive,
+                    timeout_secs=timeout,
+                ),
+                task_id=self._task_id,
+            ),
+        )
+        async with TaskContext() as tc:
+            tc.create_task(self._consume_watch_output(resp.exec_id))
+
+            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"):
+                        try:
+                            event_json = json.loads(buffer.strip().decode())
+                            event = FileWatchEvent(
+                                type=FileWatchEventType(event_json["event_type"]),
+                                paths=event_json["paths"],
+                            )
+                            if not filter or event.type in filter:
+                                yield event
+                        except (json.JSONDecodeError, KeyError, ValueError):
+                            # skip invalid events
+                            pass
+                        buffer = b""
+                else:
+                    await asyncio.sleep(0.1)
+
+    async def _close(self) -> None:
+        # Buffer is flushed by the runner on close
+        resp = await retry_transient_errors(
+            self._client.stub.ContainerFilesystemExec,
+            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")
+
+    async def __aenter__(self) -> "_FileIO":
+        return self
+
+    async def __aexit__(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)