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

modal/_ipython.py

@@ -1,21 +1,11 @@
 # Copyright Modal Labs 2022
 import sys
-import warnings
-
-ipy_outstream = None
-try:
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore")
-        import ipykernel.iostream
-
-    ipy_outstream = ipykernel.iostream.OutStream
-except ImportError:
-    pass
 
 
 def is_notebook(stdout=None):
-    if ipy_outstream is None:
+    ipykernel_iostream = sys.modules.get("ipykernel.iostream")
+    if ipykernel_iostream is None:
         return False
     if stdout is None:
         stdout = sys.stdout
-    return isinstance(stdout, ipy_outstream)
+    return isinstance(stdout, ipykernel_iostream.OutStream)

modal/_location.py

@@ -1,33 +1,40 @@
 # Copyright Modal Labs 2022
 from enum import Enum
 
-from modal_proto import api_pb2
+import modal_proto.api_pb2
 
 from .exception import InvalidError
 
 
 class CloudProvider(Enum):
-    AWS = api_pb2.CLOUD_PROVIDER_AWS
-    GCP = api_pb2.CLOUD_PROVIDER_GCP
-    AUTO = api_pb2.CLOUD_PROVIDER_AUTO
-    OCI = api_pb2.CLOUD_PROVIDER_OCI
+    AWS = modal_proto.api_pb2.CLOUD_PROVIDER_AWS
+    GCP = modal_proto.api_pb2.CLOUD_PROVIDER_GCP
+    AUTO = modal_proto.api_pb2.CLOUD_PROVIDER_AUTO
+    OCI = modal_proto.api_pb2.CLOUD_PROVIDER_OCI
 
 
-def parse_cloud_provider(value: str) -> "api_pb2.CloudProvider.V":
+def parse_cloud_provider(value: str) -> "modal_proto.api_pb2.CloudProvider.V":
     try:
         cloud_provider = CloudProvider[value.upper()]
     except KeyError:
+        # provider's int identifier may be directly specified
+        try:
+            return int(value)  # type: ignore
+        except ValueError:
+            pass
+
         raise InvalidError(
-            f"Invalid cloud provider: {value}. Value must be one of {[x.name.lower() for x in CloudProvider]} (case-insensitive)."
+            f"Invalid cloud provider: {value}. "
+            f"Value must be one of {[x.name.lower() for x in CloudProvider]} (case-insensitive)."
         )
 
     return cloud_provider.value
 
 
-def display_location(cloud_provider: "api_pb2.CloudProvider.V") -> str:
-    if cloud_provider == api_pb2.CLOUD_PROVIDER_GCP:
+def display_location(cloud_provider: "modal_proto.api_pb2.CloudProvider.V") -> str:
+    if cloud_provider == modal_proto.api_pb2.CLOUD_PROVIDER_GCP:
         return "GCP (us-east1)"
-    elif cloud_provider == api_pb2.CLOUD_PROVIDER_AWS:
+    elif cloud_provider == modal_proto.api_pb2.CLOUD_PROVIDER_AWS:
         return "AWS (us-east-1)"
     else:
         return ""

modal/_output.py

@@ -7,9 +7,11 @@ import functools
 import io
 import platform
 import re
+import socket
 import sys
+from collections.abc import Generator
 from datetime import timedelta
-from typing import Callable, Dict, Optional, Tuple
+from typing import Callable, ClassVar
 
 from grpclib.exceptions import GRPCError, StreamTerminatedError
 from rich.console import Console, Group, RenderableType
@@ -32,7 +34,7 @@ from rich.text import Text
 
 from modal_proto import api_pb2
 
-from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors, unary_stream
+from ._utils.grpc_utils import RETRYABLE_GRPC_STATUS_CODES, retry_transient_errors
 from ._utils.shell_utils import stream_from_stdin
 from .client import _Client
 from .config import logger
@@ -60,25 +62,6 @@ class FunctionQueuingColumn(ProgressColumn):
         return Text(str(delta), style="progress.elapsed")
 
 
-def step_progress(text: str = "") -> Spinner:
-    """Returns the element to be rendered when a step is in progress."""
-    return Spinner(default_spinner, text, style="blue")
-
-
-def step_progress_update(spinner: Spinner, message: str):
-    spinner.update(text=message)
-
-
-def step_completed(message: str, is_substep: bool = False) -> RenderableType:
-    """Returns the element to be rendered when a step is completed."""
-
-    STEP_COMPLETED = "[green]✓[/green]"
-    SUBSTEP_COMPLETED = "🔨"
-
-    symbol = SUBSTEP_COMPLETED if is_substep else STEP_COMPLETED
-    return f"{symbol} {message}"
-
-
 def download_progress_bar() -> Progress:
     """
     Returns a progress bar suitable for showing file download progress.
@@ -139,24 +122,28 @@ class LineBufferedOutput(io.StringIO):
 
 
 class OutputManager:
-    _visible_progress: bool
+    _instance: ClassVar[OutputManager | None] = None
+
     _console: Console
-    _task_states: Dict[str, int]
-    _task_progress_items: Dict[Tuple[str, int], TaskID]
-    _current_render_group: Optional[Group]
-    _function_progress: Optional[Progress]
-    _function_queueing_progress: Optional[Progress]
-    _snapshot_progress: Optional[Progress]
-    _line_buffers: Dict[int, LineBufferedOutput]
+    _task_states: dict[str, int]
+    _task_progress_items: dict[tuple[str, int], TaskID]
+    _current_render_group: Group | None
+    _function_progress: Progress | None
+    _function_queueing_progress: Progress | None
+    _snapshot_progress: Progress | None
+    _line_buffers: dict[int, LineBufferedOutput]
     _status_spinner: Spinner
-    _app_page_url: Optional[str]
+    _app_page_url: str | None
     _show_image_logs: bool
-    _status_spinner_live: Optional[Live]
-
-    def __init__(self, stdout: io.TextIOWrapper, show_progress: bool, status_spinner_text: str = "Running app..."):
-        self.stdout = stdout or sys.stdout
-
-        self._visible_progress = show_progress
+    _status_spinner_live: Live | None
+
+    def __init__(
+        self,
+        *,
+        stdout: io.TextIOWrapper | None = None,
+        status_spinner_text: str = "Running app...",
+    ):
+        self._stdout = stdout or sys.stdout
         self._console = Console(file=stdout, highlight=False)
         self._task_states = {}
         self._task_progress_items = {}
@@ -165,18 +152,47 @@ class OutputManager:
         self._function_queueing_progress = None
         self._snapshot_progress = None
         self._line_buffers = {}
-        self._status_spinner = step_progress(status_spinner_text)
+        self._status_spinner = OutputManager.step_progress(status_spinner_text)
         self._app_page_url = None
         self._show_image_logs = False
+        self._status_spinner_live = None
 
-    def print_if_visible(self, renderable) -> None:
-        if self._visible_progress:
-            self._console.print(renderable)
+    @classmethod
+    def disable(cls):
+        cls._instance.flush_lines()
+        if cls._instance._status_spinner_live:
+            cls._instance._status_spinner_live.stop()
+        cls._instance = None
 
-    def ctx_if_visible(self, context_mgr):
-        if self._visible_progress:
-            return context_mgr
-        return contextlib.nullcontext()
+    @classmethod
+    def get(cls) -> OutputManager | None:
+        return cls._instance
+
+    @classmethod
+    @contextlib.contextmanager
+    def enable_output(cls, show_progress: bool = True) -> Generator[None]:
+        if show_progress:
+            cls._instance = OutputManager()
+        try:
+            yield
+        finally:
+            cls._instance = None
+
+    @staticmethod
+    def step_progress(text: str = "") -> Spinner:
+        """Returns the element to be rendered when a step is in progress."""
+        return Spinner(default_spinner, text, style="blue")
+
+    @staticmethod
+    def step_completed(message: str) -> RenderableType:
+        return f"[green]✓[/green] {message}"
+
+    @staticmethod
+    def substep_completed(message: str) -> RenderableType:
+        return f"🔨 {message}"
+
+    def print(self, renderable) -> None:
+        self._console.print(renderable)
 
     def make_live(self, renderable: RenderableType) -> Live:
         """Creates a customized `rich.Live` instance with the given renderable. The renderable
@@ -237,7 +253,7 @@ class OutputManager:
                 self._current_render_group.renderables.append(self._function_queueing_progress)
         return self._function_queueing_progress
 
-    def function_progress_callback(self, tag: str, total: Optional[int]) -> Callable[[int, int], None]:
+    def function_progress_callback(self, tag: str, total: int | None) -> Callable[[int, int], None]:
         """Adds a task to the current function_progress instance, and returns a callback
         to update task progress with new completed and total counts."""
 
@@ -292,7 +308,7 @@ class OutputManager:
         message = f"[blue]{message}[/blue] [grey70]View app at [underline]{self._app_page_url}[/underline][/grey70]"
 
         # Set the new message
-        step_progress_update(self._status_spinner, message)
+        self._status_spinner.update(text=message)
 
     def update_snapshot_progress(self, image_id: str, task_progress: api_pb2.TaskProgress) -> None:
         # TODO(erikbern): move this to sit on the resolver object, mostly
@@ -315,7 +331,7 @@ class OutputManager:
             pass
 
     def update_queueing_progress(
-        self, *, function_id: str, completed: int, total: Optional[int], description: Optional[str]
+        self, *, function_id: str, completed: int, total: int | None, description: str | None
     ) -> None:
         """Handle queueing updates, ignoring completion updates for functions that have no queue progress bar."""
         task_key = (function_id, api_pb2.FUNCTION_QUEUED)
@@ -335,33 +351,11 @@ class OutputManager:
             self._task_progress_items[task_key] = progress_task_id
 
     async def put_log_content(self, log: api_pb2.TaskLogs):
-        if self._visible_progress:
-            stream = self._line_buffers.get(log.file_descriptor)
-            if stream is None:
-                stream = LineBufferedOutput(functools.partial(self._print_log, log.file_descriptor))
-                self._line_buffers[log.file_descriptor] = stream
-            stream.write(log.data)
-        elif hasattr(self.stdout, "buffer"):
-            # If we're not showing progress, there's no need to buffer lines,
-            # because the progress spinner can't interfere with output.
-
-            data = log.data.encode("utf-8")
-            written = 0
-            n_retries = 0
-            while written < len(data):
-                try:
-                    written += self.stdout.buffer.write(data[written:])
-                    self.stdout.flush()
-                except BlockingIOError:
-                    if n_retries >= 5:
-                        raise
-                    n_retries += 1
-                    await asyncio.sleep(0.1)
-        else:
-            # `stdout` isn't always buffered (e.g. %%capture in Jupyter notebooks redirects it to
-            # io.StringIO).
-            self.stdout.write(log.data)
-            self.stdout.flush()
+        stream = self._line_buffers.get(log.file_descriptor)
+        if stream is None:
+            stream = LineBufferedOutput(functools.partial(self._print_log, log.file_descriptor))
+            self._line_buffers[log.file_descriptor] = stream
+        stream.write(log.data)
 
     def flush_lines(self):
         for stream in self._line_buffers.values():
@@ -370,12 +364,123 @@ class OutputManager:
     @contextlib.contextmanager
     def show_status_spinner(self):
         self._status_spinner_live = self.make_live(self._status_spinner)
-        with self.ctx_if_visible(self._status_spinner_live):
+        with self._status_spinner_live:
             yield
 
-    def hide_status_spinner(self):
-        if self._status_spinner_live:
-            self._status_spinner_live.stop()
+
+class ProgressHandler:
+    live: Live
+    _type: str
+    _spinner: Spinner
+    _overall_progress: Progress
+    _download_progress: Progress
+    _overall_progress_task_id: TaskID
+    _total_tasks: int
+    _completed_tasks: int
+
+    def __init__(self, type: str, console: Console):
+        self._type = type
+
+        if self._type == "download":
+            title = "Downloading file(s) to local..."
+        elif self._type == "upload":
+            title = "Uploading file(s) to volume..."
+        else:
+            raise NotImplementedError(f"Progress handler of type: `{type}` not yet implemented")
+
+        self._spinner = OutputManager.step_progress(title)
+
+        self._overall_progress = Progress(
+            TextColumn(f"[bold white]{title}", justify="right"),
+            TimeElapsedColumn(),
+            BarColumn(bar_width=None),
+            TextColumn("[bold white]{task.description}"),
+            transient=True,
+            console=console,
+        )
+        self._download_progress = Progress(
+            TextColumn("[bold white]{task.fields[path]}", justify="right"),
+            BarColumn(bar_width=None),
+            "[progress.percentage]{task.percentage:>3.1f}%",
+            "•",
+            DownloadColumn(),
+            "•",
+            TransferSpeedColumn(),
+            "•",
+            TimeRemainingColumn(),
+            transient=True,
+            console=console,
+        )
+
+        self.live = Live(
+            Group(self._spinner, self._overall_progress, self._download_progress), transient=True, refresh_per_second=4
+        )
+
+        self._overall_progress_task_id = self._overall_progress.add_task(".", start=True)
+        self._total_tasks = 0
+        self._completed_tasks = 0
+
+    def _add_sub_task(self, name: str, size: float) -> TaskID:
+        task_id = self._download_progress.add_task(self._type, path=name, start=True, total=size)
+        self._total_tasks += 1
+        self._overall_progress.update(self._overall_progress_task_id, total=self._total_tasks)
+        return task_id
+
+    def _reset_sub_task(self, task_id: TaskID):
+        self._download_progress.reset(task_id)
+
+    def _complete_progress(self):
+        # TODO: we could probably implement some callback progression from the server
+        # to get progress reports for the post processing too
+        # so we don't have to just spin here
+        self._overall_progress.remove_task(self._overall_progress_task_id)
+        self._spinner.update(text="Post processing...")
+
+    def _complete_sub_task(self, task_id: TaskID):
+        self._completed_tasks += 1
+        self._download_progress.remove_task(task_id)
+        self._overall_progress.update(
+            self._overall_progress_task_id,
+            advance=1,
+            description=f"({self._completed_tasks} out of {self._total_tasks} files completed)",
+        )
+
+    def _advance_sub_task(self, task_id: TaskID, advance: float):
+        self._download_progress.update(task_id, advance=advance)
+
+    def progress(
+        self,
+        task_id: TaskID | None = None,
+        advance: float | None = None,
+        name: str | None = None,
+        size: float | None = None,
+        reset: bool | None = False,
+        complete: bool | None = False,
+    ) -> TaskID | None:
+        try:
+            if task_id is not None:
+                if reset:
+                    return self._reset_sub_task(task_id)
+                elif complete:
+                    return self._complete_sub_task(task_id)
+                elif advance is not None:
+                    return self._advance_sub_task(task_id, advance)
+            elif name is not None and size is not None:
+                return self._add_sub_task(name, size)
+            elif complete:
+                return self._complete_progress()
+        except Exception as exc:
+            # Liberal exception handling to avoid crashing downloads and uploads.
+            logger.error(f"failed progress update: {exc}")
+        raise NotImplementedError(
+            "Unknown action to take with args: "
+            + f"name={name} "
+            + f"size={size} "
+            + f"task_id={task_id} "
+            + f"advance={advance} "
+            + f"reset={reset} "
+            + f"complete={complete} "
+        )
 
 
 async def stream_pty_shell_input(client: _Client, exec_id: str, finish_event: asyncio.Event):
@@ -396,10 +501,42 @@ async def stream_pty_shell_input(client: _Client, exec_id: str, finish_event: as
         await finish_event.wait()
 
 
-async def get_app_logs_loop(app_id: str, client: _Client, output_mgr: OutputManager):
+async def put_pty_content(log: api_pb2.TaskLogs, stdout):
+    if hasattr(stdout, "buffer"):
+        # If we're not showing progress, there's no need to buffer lines,
+        # because the progress spinner can't interfere with output.
+
+        data = log.data.encode("utf-8")
+        written = 0
+        n_retries = 0
+        while written < len(data):
+            try:
+                written += stdout.buffer.write(data[written:])
+                stdout.flush()
+            except BlockingIOError:
+                if n_retries >= 5:
+                    raise
+                n_retries += 1
+                await asyncio.sleep(0.1)
+    else:
+        # `stdout` isn't always buffered (e.g. %%capture in Jupyter notebooks redirects it to
+        # io.StringIO).
+        stdout.write(log.data)
+        stdout.flush()
+
+
+async def get_app_logs_loop(
+    client: _Client,
+    output_mgr: OutputManager,
+    app_id: str | None = None,
+    task_id: str | None = None,
+    app_logs_url: str | None = None,
+):
     last_log_batch_entry_id = ""
-    pty_shell_finish_event: Optional[asyncio.Event] = None
-    pty_shell_task_id: Optional[str] = None
+
+    pty_shell_stdout = None
+    pty_shell_finish_event: asyncio.Event | None = None
+    pty_shell_task_id: str | None = None
 
     async def stop_pty_shell():
         nonlocal pty_shell_finish_event
@@ -428,18 +565,23 @@ async def get_app_logs_loop(app_id: str, client: _Client, output_mgr: OutputMana
             else:  # Ensure forward-compatible with new types.
                 logger.debug(f"Received unrecognized progress type: {log.task_progress.progress_type}")
         elif log.data:
-            await output_mgr.put_log_content(log)
+            if pty_shell_finish_event:
+                await put_pty_content(log, pty_shell_stdout)
+            else:
+                await output_mgr.put_log_content(log)
 
     async def _get_logs():
-        nonlocal last_log_batch_entry_id, pty_shell_finish_event, pty_shell_task_id
+        nonlocal last_log_batch_entry_id
+        nonlocal pty_shell_stdout, pty_shell_finish_event, pty_shell_task_id
 
         request = api_pb2.AppGetLogsRequest(
-            app_id=app_id,
+            app_id=app_id or "",
+            task_id=task_id or "",
             timeout=55,
             last_entry_id=last_log_batch_entry_id,
         )
         log_batch: api_pb2.TaskLogsBatch
-        async for log_batch in unary_stream(client.stub.AppGetLogs, request):
+        async for log_batch in client.stub.AppGetLogs.unary_stream(request):
             if log_batch.entry_id:
                 # log_batch entry_id is empty for fd="server" messages from AppGetLogs
                 last_log_batch_entry_id = log_batch.entry_id
@@ -456,14 +598,15 @@ async def get_app_logs_loop(app_id: str, client: _Client, output_mgr: OutputMana
                 # statically and dynamically built images.
                 pass
             elif log_batch.pty_exec_id:
+                # This corresponds to the `modal run -i` use case where a breakpoint
+                # triggers and the task drops into an interactive PTY mode
                 if pty_shell_finish_event:
                     print("ERROR: concurrent PTY shells are not supported.")
                 else:
-                    output_mgr.flush_lines()
-                    output_mgr.hide_status_spinner()
-                    output_mgr._visible_progress = False
+                    pty_shell_stdout = output_mgr._stdout
                     pty_shell_finish_event = asyncio.Event()
                     pty_shell_task_id = log_batch.task_id
+                    output_mgr.disable()
                     asyncio.create_task(stream_pty_shell_input(client, log_batch.pty_exec_id, pty_shell_finish_event))
             else:
                 for log in log_batch.items:
@@ -477,14 +620,7 @@ async def get_app_logs_loop(app_id: str, client: _Client, output_mgr: OutputMana
     while True:
         try:
             await _get_logs()
-        except asyncio.CancelledError:
-            # TODO: this should come from the backend maybe
-            app_logs_url = f"https://modal.com/logs/{app_id}"
-            output_mgr.print_if_visible(
-                f"[red]Timed out waiting for logs. [grey70]View logs at [underline]{app_logs_url}[/underline] for remaining output.[/grey70]"
-            )
-            raise
-        except (GRPCError, StreamTerminatedError) as exc:
+        except (GRPCError, StreamTerminatedError, socket.gaierror, AttributeError) as exc:
             if isinstance(exc, GRPCError):
                 if exc.status in RETRYABLE_GRPC_STATUS_CODES:
                     # Try again if we had a temporary connection drop,
@@ -494,10 +630,18 @@ async def get_app_logs_loop(app_id: str, client: _Client, output_mgr: OutputMana
             elif isinstance(exc, StreamTerminatedError):
                 logger.debug("Stream closed. Retrying ...")
                 continue
+            elif isinstance(exc, socket.gaierror):
+                logger.debug("Lost connection. Retrying ...")
+                continue
+            elif isinstance(exc, AttributeError):
+                if "_write_appdata" in str(exc):
+                    # Happens after losing connection
+                    # StreamTerminatedError are not properly raised in grpclib<=0.4.7
+                    # fixed in https://github.com/vmagamedov/grpclib/issues/185
+                    # TODO: update to newer version (>=0.4.8) once stable
+                    logger.debug("Lost connection. Retrying ...")
+                    continue
             raise
-        except Exception as exc:
-            logger.exception(f"Failed to fetch logs: {exc}")
-            await asyncio.sleep(1)
 
         if last_log_batch_entry_id is None:
             break

modal/_pty.py

@@ -2,12 +2,12 @@
 import contextlib
 import os
 import sys
-from typing import Optional, Tuple
+from typing import Optional
 
 from modal_proto import api_pb2
 
 
-def get_winsz(fd) -> Tuple[Optional[int], Optional[int]]:
+def get_winsz(fd) -> tuple[Optional[int], Optional[int]]:
     try:
         import fcntl
         import struct