modal 1.0.4.dev12__py3-none-any.whl → 1.0.5__py3-none-any.whl

modal/_tunnel.pyi

@@ -4,34 +4,402 @@ import typing
 import typing_extensions
 
 class Tunnel:
+    """A port forwarded from within a running Modal container. Created by `modal.forward()`.
+
+    **Important:** This is an experimental API which may change in the future.
+    """
+
     host: str
     port: int
     unencrypted_host: str
     unencrypted_port: int
 
     @property
-    def url(self) -> str: ...
+    def url(self) -> str:
+        """Get the public HTTPS URL of the forwarded port."""
+        ...
+
     @property
-    def tls_socket(self) -> tuple[str, int]: ...
+    def tls_socket(self) -> tuple[str, int]:
+        """Get the public TLS socket as a (host, port) tuple."""
+        ...
+
     @property
-    def tcp_socket(self) -> tuple[str, int]: ...
-    def __init__(self, host: str, port: int, unencrypted_host: str, unencrypted_port: int) -> None: ...
-    def __repr__(self): ...
-    def __eq__(self, other): ...
-    def __setattr__(self, name, value): ...
-    def __delattr__(self, name): ...
-    def __hash__(self): ...
+    def tcp_socket(self) -> tuple[str, int]:
+        """Get the public TCP socket as a (host, port) tuple."""
+        ...
+
+    def __init__(self, host: str, port: int, unencrypted_host: str, unencrypted_port: int) -> None:
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def __repr__(self):
+        """Return repr(self)."""
+        ...
+
+    def __eq__(self, other):
+        """Return self==value."""
+        ...
+
+    def __setattr__(self, name, value):
+        """Implement setattr(self, name, value)."""
+        ...
+
+    def __delattr__(self, name):
+        """Implement delattr(self, name)."""
+        ...
+
+    def __hash__(self):
+        """Return hash(self)."""
+        ...
 
 def _forward(
     port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client._Client] = None
-) -> typing.AsyncContextManager[Tunnel]: ...
+) -> typing.AsyncContextManager[Tunnel]:
+    '''Expose a port publicly from inside a running Modal container, with TLS.
+
+    If `unencrypted` is set, this also exposes the TCP socket without encryption on a random port
+    number. This can be used to SSH into a container (see example below). Note that it is on the public Internet, so
+    make sure you are using a secure protocol over TCP.
+
+    **Important:** This is an experimental API which may change in the future.
+
+    **Usage:**
+
+    ```python notest
+    import modal
+    from flask import Flask
+
+    app = modal.App(image=modal.Image.debian_slim().pip_install("Flask"))
+    flask_app = Flask(__name__)
+
+
+    @flask_app.route("/")
+    def hello_world():
+        return "Hello, World!"
+
+
+    @app.function()
+    def run_app():
+        # Start a web server inside the container at port 8000. `modal.forward(8000)` lets us
+        # expose that port to the world at a random HTTPS URL.
+        with modal.forward(8000) as tunnel:
+            print("Server listening at", tunnel.url)
+            flask_app.run("0.0.0.0", 8000)
+
+        # When the context manager exits, the port is no longer exposed.
+    ```
+
+    **Raw TCP usage:**
+
+    ```python
+    import socket
+    import threading
+
+    import modal
+
+
+    def run_echo_server(port: int):
+        """Run a TCP echo server listening on the given port."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.bind(("0.0.0.0", port))
+        sock.listen(1)
+
+        while True:
+            conn, addr = sock.accept()
+            print("Connection from:", addr)
+
+            # Start a new thread to handle the connection
+            def handle(conn):
+                with conn:
+                    while True:
+                        data = conn.recv(1024)
+                        if not data:
+                            break
+                        conn.sendall(data)
+
+            threading.Thread(target=handle, args=(conn,)).start()
+
+
+    app = modal.App()
+
+
+    @app.function()
+    def tcp_tunnel():
+        # This exposes port 8000 to public Internet traffic over TCP.
+        with modal.forward(8000, unencrypted=True) as tunnel:
+            # You can connect to this TCP socket from outside the container, for example, using `nc`:
+            #  nc <HOST> <PORT>
+            print("TCP tunnel listening at:", tunnel.tcp_socket)
+            run_echo_server(8000)
+    ```
+
+    **SSH example:**
+    This assumes you have a rsa keypair in `~/.ssh/id_rsa{.pub}`, this is a bare-bones example
+    letting you SSH into a Modal container.
+
+    ```python
+    import subprocess
+    import time
+
+    import modal
+
+    app = modal.App()
+    image = (
+        modal.Image.debian_slim()
+        .apt_install("openssh-server")
+        .run_commands("mkdir /run/sshd")
+        .add_local_file("~/.ssh/id_rsa.pub", "/root/.ssh/authorized_keys", copy=True)
+    )
+
+
+    @app.function(image=image, timeout=3600)
+    def some_function():
+        subprocess.Popen(["/usr/sbin/sshd", "-D", "-e"])
+        with modal.forward(port=22, unencrypted=True) as tunnel:
+            hostname, port = tunnel.tcp_socket
+            connection_cmd = f'ssh -p {port} root@{hostname}'
+            print(f"ssh into container using: {connection_cmd}")
+            time.sleep(3600)  # keep alive for 1 hour or until killed
+    ```
+
+    If you intend to use this more generally, a suggestion is to put the subprocess and port
+    forwarding code in an `@enter` lifecycle method of an @app.cls, to only make a single
+    ssh server and port for each container (and not one for each input to the function).
+    '''
+    ...
 
 class __forward_spec(typing_extensions.Protocol):
     def __call__(
         self, /, port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client.Client] = None
-    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[Tunnel]: ...
+    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[Tunnel]:
+        '''Expose a port publicly from inside a running Modal container, with TLS.
+
+        If `unencrypted` is set, this also exposes the TCP socket without encryption on a random port
+        number. This can be used to SSH into a container (see example below). Note that it is on the public Internet, so
+        make sure you are using a secure protocol over TCP.
+
+        **Important:** This is an experimental API which may change in the future.
+
+        **Usage:**
+
+        ```python notest
+        import modal
+        from flask import Flask
+
+        app = modal.App(image=modal.Image.debian_slim().pip_install("Flask"))
+        flask_app = Flask(__name__)
+
+
+        @flask_app.route("/")
+        def hello_world():
+            return "Hello, World!"
+
+
+        @app.function()
+        def run_app():
+            # Start a web server inside the container at port 8000. `modal.forward(8000)` lets us
+            # expose that port to the world at a random HTTPS URL.
+            with modal.forward(8000) as tunnel:
+                print("Server listening at", tunnel.url)
+                flask_app.run("0.0.0.0", 8000)
+
+            # When the context manager exits, the port is no longer exposed.
+        ```
+
+        **Raw TCP usage:**
+
+        ```python
+        import socket
+        import threading
+
+        import modal
+
+
+        def run_echo_server(port: int):
+            """Run a TCP echo server listening on the given port."""
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.bind(("0.0.0.0", port))
+            sock.listen(1)
+
+            while True:
+                conn, addr = sock.accept()
+                print("Connection from:", addr)
+
+                # Start a new thread to handle the connection
+                def handle(conn):
+                    with conn:
+                        while True:
+                            data = conn.recv(1024)
+                            if not data:
+                                break
+                            conn.sendall(data)
+
+                threading.Thread(target=handle, args=(conn,)).start()
+
+
+        app = modal.App()
+
+
+        @app.function()
+        def tcp_tunnel():
+            # This exposes port 8000 to public Internet traffic over TCP.
+            with modal.forward(8000, unencrypted=True) as tunnel:
+                # You can connect to this TCP socket from outside the container, for example, using `nc`:
+                #  nc <HOST> <PORT>
+                print("TCP tunnel listening at:", tunnel.tcp_socket)
+                run_echo_server(8000)
+        ```
+
+        **SSH example:**
+        This assumes you have a rsa keypair in `~/.ssh/id_rsa{.pub}`, this is a bare-bones example
+        letting you SSH into a Modal container.
+
+        ```python
+        import subprocess
+        import time
+
+        import modal
+
+        app = modal.App()
+        image = (
+            modal.Image.debian_slim()
+            .apt_install("openssh-server")
+            .run_commands("mkdir /run/sshd")
+            .add_local_file("~/.ssh/id_rsa.pub", "/root/.ssh/authorized_keys", copy=True)
+        )
+
+
+        @app.function(image=image, timeout=3600)
+        def some_function():
+            subprocess.Popen(["/usr/sbin/sshd", "-D", "-e"])
+            with modal.forward(port=22, unencrypted=True) as tunnel:
+                hostname, port = tunnel.tcp_socket
+                connection_cmd = f'ssh -p {port} root@{hostname}'
+                print(f"ssh into container using: {connection_cmd}")
+                time.sleep(3600)  # keep alive for 1 hour or until killed
+        ```
+
+        If you intend to use this more generally, a suggestion is to put the subprocess and port
+        forwarding code in an `@enter` lifecycle method of an @app.cls, to only make a single
+        ssh server and port for each container (and not one for each input to the function).
+        '''
+        ...
+
     def aio(
         self, /, port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client.Client] = None
-    ) -> typing.AsyncContextManager[Tunnel]: ...
+    ) -> typing.AsyncContextManager[Tunnel]:
+        '''Expose a port publicly from inside a running Modal container, with TLS.
+
+        If `unencrypted` is set, this also exposes the TCP socket without encryption on a random port
+        number. This can be used to SSH into a container (see example below). Note that it is on the public Internet, so
+        make sure you are using a secure protocol over TCP.
+
+        **Important:** This is an experimental API which may change in the future.
+
+        **Usage:**
+
+        ```python notest
+        import modal
+        from flask import Flask
+
+        app = modal.App(image=modal.Image.debian_slim().pip_install("Flask"))
+        flask_app = Flask(__name__)
+
+
+        @flask_app.route("/")
+        def hello_world():
+            return "Hello, World!"
+
+
+        @app.function()
+        def run_app():
+            # Start a web server inside the container at port 8000. `modal.forward(8000)` lets us
+            # expose that port to the world at a random HTTPS URL.
+            with modal.forward(8000) as tunnel:
+                print("Server listening at", tunnel.url)
+                flask_app.run("0.0.0.0", 8000)
+
+            # When the context manager exits, the port is no longer exposed.
+        ```
+
+        **Raw TCP usage:**
+
+        ```python
+        import socket
+        import threading
+
+        import modal
+
+
+        def run_echo_server(port: int):
+            """Run a TCP echo server listening on the given port."""
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.bind(("0.0.0.0", port))
+            sock.listen(1)
+
+            while True:
+                conn, addr = sock.accept()
+                print("Connection from:", addr)
+
+                # Start a new thread to handle the connection
+                def handle(conn):
+                    with conn:
+                        while True:
+                            data = conn.recv(1024)
+                            if not data:
+                                break
+                            conn.sendall(data)
+
+                threading.Thread(target=handle, args=(conn,)).start()
+
+
+        app = modal.App()
+
+
+        @app.function()
+        def tcp_tunnel():
+            # This exposes port 8000 to public Internet traffic over TCP.
+            with modal.forward(8000, unencrypted=True) as tunnel:
+                # You can connect to this TCP socket from outside the container, for example, using `nc`:
+                #  nc <HOST> <PORT>
+                print("TCP tunnel listening at:", tunnel.tcp_socket)
+                run_echo_server(8000)
+        ```
+
+        **SSH example:**
+        This assumes you have a rsa keypair in `~/.ssh/id_rsa{.pub}`, this is a bare-bones example
+        letting you SSH into a Modal container.
+
+        ```python
+        import subprocess
+        import time
+
+        import modal
+
+        app = modal.App()
+        image = (
+            modal.Image.debian_slim()
+            .apt_install("openssh-server")
+            .run_commands("mkdir /run/sshd")
+            .add_local_file("~/.ssh/id_rsa.pub", "/root/.ssh/authorized_keys", copy=True)
+        )
+
+
+        @app.function(image=image, timeout=3600)
+        def some_function():
+            subprocess.Popen(["/usr/sbin/sshd", "-D", "-e"])
+            with modal.forward(port=22, unencrypted=True) as tunnel:
+                hostname, port = tunnel.tcp_socket
+                connection_cmd = f'ssh -p {port} root@{hostname}'
+                print(f"ssh into container using: {connection_cmd}")
+                time.sleep(3600)  # keep alive for 1 hour or until killed
+        ```
+
+        If you intend to use this more generally, a suggestion is to put the subprocess and port
+        forwarding code in an `@enter` lifecycle method of an @app.cls, to only make a single
+        ssh server and port for each container (and not one for each input to the function).
+        '''
+        ...
 
 forward: __forward_spec

modal/_utils/async_utils.py

@@ -396,7 +396,7 @@ class _WarnIfGeneratorIsNotConsumed:
         return await self.gen.aclose()
 
 
-synchronize_api(_WarnIfGeneratorIsNotConsumed)
+_BlockingWarnIfGeneratorIsNotConsumed = synchronize_api(_WarnIfGeneratorIsNotConsumed)
 
 
 class _WarnIfNonWrappedGeneratorIsNotConsumed(_WarnIfGeneratorIsNotConsumed):

modal/_utils/blob_utils.py

@@ -4,6 +4,7 @@ import dataclasses
 import hashlib
 import os
 import platform
+import random
 import time
 from collections.abc import AsyncIterator
 from contextlib import AbstractContextManager, contextmanager
@@ -37,12 +38,15 @@ if TYPE_CHECKING:
 # Max size for function inputs and outputs.
 MAX_OBJECT_SIZE_BYTES = 2 * 1024 * 1024  # 2 MiB
 
+# Max size for async function inputs and outputs.
+MAX_ASYNC_OBJECT_SIZE_BYTES = 8 * 1024  # 8 KiB
+
 #  If a file is LARGE_FILE_LIMIT bytes or larger, it's uploaded to blob store (s3) instead of going through grpc
 #  It will also make sure to chunk the hash calculation to avoid reading the entire file into memory
 LARGE_FILE_LIMIT = 4 * 1024 * 1024  # 4 MiB
 
 # Max parallelism during map calls
-BLOB_MAX_PARALLELISM = 10
+BLOB_MAX_PARALLELISM = 20
 
 # read ~16MiB chunks by default
 DEFAULT_SEGMENT_CHUNK_SIZE = 2**24
@@ -55,6 +59,8 @@ MULTIPART_UPLOAD_THRESHOLD = 1024**3
 # For block based storage like volumefs2: the size of a block
 BLOCK_SIZE: int = 8 * 1024 * 1024
 
+HEALTHY_R2_UPLOAD_PERCENTAGE = 0.95
+
 
 @retry(n_attempts=5, base_delay=0.5, timeout=None)
 async def _upload_to_s3_url(
@@ -182,6 +188,22 @@ def get_content_length(data: BinaryIO) -> int:
     return content_length - pos
 
 
+async def _blob_upload_with_fallback(items, blob_ids, callback):
+    for idx, (item, blob_id) in enumerate(zip(items, blob_ids)):
+        # We want to default to R2 95% of the time and S3 5% of the time.
+        # To ensure the failure path is continuously exercised.
+        if idx == 0 and len(items) > 1 and random.random() > HEALTHY_R2_UPLOAD_PERCENTAGE:
+            continue
+        try:
+            await callback(item)
+            return blob_id
+        except Exception as _:
+            # Ignore all errors except the last one, since we're out of fallback options.
+            if idx == len(items) - 1:
+                raise
+    raise ExecutionError("Failed to upload blob")
+
+
 async def _blob_upload(
     upload_hashes: UploadHashes, data: Union[bytes, BinaryIO], stub, progress_report_cb: Optional[Callable] = None
 ) -> str:
@@ -197,17 +219,23 @@ async def _blob_upload(
     )
     resp = await retry_transient_errors(stub.BlobCreate, req)
 
-    blob_id = resp.blob_id
+    if resp.WhichOneof("upload_types_oneof") == "multiparts":
+
+        async def upload_multipart_upload(part):
+            return await perform_multipart_upload(
+                data,
+                content_length=content_length,
+                max_part_size=part.part_length,
+                part_urls=part.upload_urls,
+                completion_url=part.completion_url,
+                upload_chunk_size=DEFAULT_SEGMENT_CHUNK_SIZE,
+                progress_report_cb=progress_report_cb,
+            )
 
-    if resp.WhichOneof("upload_type_oneof") == "multipart":
-        await perform_multipart_upload(
-            data,
-            content_length=content_length,
-            max_part_size=resp.multipart.part_length,
-            part_urls=resp.multipart.upload_urls,
-            completion_url=resp.multipart.completion_url,
-            upload_chunk_size=DEFAULT_SEGMENT_CHUNK_SIZE,
-            progress_report_cb=progress_report_cb,
+        blob_id = await _blob_upload_with_fallback(
+            resp.multiparts.items,
+            resp.blob_ids,
+            upload_multipart_upload,
         )
     else:
         from .bytes_io_segment_payload import BytesIOSegmentPayload
@@ -215,11 +243,19 @@ async def _blob_upload(
         payload = BytesIOSegmentPayload(
             data, segment_start=0, segment_length=content_length, progress_report_cb=progress_report_cb
         )
-        await _upload_to_s3_url(
-            resp.upload_url,
-            payload,
-            # for single part uploads, we use server side md5 checksums
-            content_md5_b64=upload_hashes.md5_base64,
+
+        async def upload_to_s3_url(url):
+            return await _upload_to_s3_url(
+                url,
+                payload,
+                # for single part uploads, we use server side md5 checksums
+                content_md5_b64=upload_hashes.md5_base64,
+            )
+
+        blob_id = await _blob_upload_with_fallback(
+            resp.upload_urls.items,
+            resp.blob_ids,
+            upload_to_s3_url,
         )
 
     if progress_report_cb:
@@ -380,8 +416,10 @@ def get_file_upload_spec_from_fileobj(fp: BinaryIO, mount_filename: PurePosixPat
         mode,
     )
 
+
 _FileUploadSource2 = Callable[[], ContextManager[BinaryIO]]
 
+
 @dataclasses.dataclass
 class FileUploadSpec2:
     source: _FileUploadSource2
@@ -393,7 +431,6 @@ class FileUploadSpec2:
     mode: int  # file permission bits (last 12 bits of st_mode)
     size: int
 
-
     @staticmethod
     async def from_path(
         filename: Path,
@@ -416,7 +453,6 @@ class FileUploadSpec2:
             hash_semaphore,
         )
 
-
     @staticmethod
     async def from_fileobj(
         source_fp: Union[BinaryIO, BytesIO],
@@ -426,6 +462,7 @@ class FileUploadSpec2:
     ) -> "FileUploadSpec2":
         try:
             fileno = source_fp.fileno()
+
             def source():
                 new_fd = os.dup(fileno)
                 fp = os.fdopen(new_fd, "rb")
@@ -436,6 +473,7 @@ class FileUploadSpec2:
             # `.fileno()` not available; assume BytesIO-like type
             source_fp = cast(BytesIO, source_fp)
             buffer = source_fp.getbuffer()
+
             def source():
                 return BytesIO(buffer)
 
@@ -447,7 +485,6 @@ class FileUploadSpec2:
             hash_semaphore,
         )
 
-
     @staticmethod
     async def _create(
         source: _FileUploadSource2,

modal/_utils/function_utils.py

@@ -10,7 +10,6 @@ from typing import Any, Callable, Literal, Optional
 
 from grpclib import GRPCError
 from grpclib.exceptions import StreamTerminatedError
-from synchronicity.exceptions import UserCodeException
 
 import modal_proto
 from modal_proto import api_pb2
@@ -33,7 +32,7 @@ from ..exception import (
     RemoteError,
 )
 from ..mount import ROOT_DIR, _is_modal_path, _Mount
-from .blob_utils import MAX_OBJECT_SIZE_BYTES, blob_download, blob_upload
+from .blob_utils import MAX_ASYNC_OBJECT_SIZE_BYTES, MAX_OBJECT_SIZE_BYTES, blob_download, blob_upload
 from .grpc_utils import RETRYABLE_GRPC_STATUS_CODES
 
 
@@ -497,8 +496,9 @@ async def _process_result(result: api_pb2.GenericResult, data_format: int, stub,
                     append_modal_tb(exc, tb_dict, line_cache)
                 except Exception:
                     pass
-            uc_exc = UserCodeException(exc_with_hints(exc))
-            raise uc_exc
+
+            raise exc_with_hints(exc)
+
         raise RemoteError(result.exception)
 
     try:
@@ -511,8 +511,27 @@ async def _process_result(result: api_pb2.GenericResult, data_format: int, stub,
         ) from deser_exc
 
 
+def should_upload(
+    num_bytes: int,
+    function_call_invocation_type: Optional["api_pb2.FunctionCallInvocationType.ValueType"],
+) -> bool:
+    """
+    Determine if the input should be uploaded to blob storage.
+    """
+    return num_bytes > MAX_OBJECT_SIZE_BYTES or (
+        function_call_invocation_type == api_pb2.FUNCTION_CALL_INVOCATION_TYPE_ASYNC
+        and num_bytes > MAX_ASYNC_OBJECT_SIZE_BYTES
+    )
+
+
 async def _create_input(
-    args, kwargs, stub: ModalClientModal, *, idx: Optional[int] = None, method_name: Optional[str] = None
+    args,
+    kwargs,
+    stub: ModalClientModal,
+    *,
+    idx: Optional[int] = None,
+    method_name: Optional[str] = None,
+    function_call_invocation_type: Optional["api_pb2.FunctionCallInvocationType.ValueType"] = None,
 ) -> api_pb2.FunctionPutInputsItem:
     """Serialize function arguments and create a FunctionInput protobuf,
     uploading to blob storage if needed.
@@ -524,9 +543,8 @@ async def _create_input(
 
     args_serialized = serialize((args, kwargs))
 
-    if len(args_serialized) > MAX_OBJECT_SIZE_BYTES:
+    if should_upload(len(args_serialized), function_call_invocation_type):
         args_blob_id = await blob_upload(args_serialized, stub)
-
         return api_pb2.FunctionPutInputsItem(
             input=api_pb2.FunctionInput(
                 args_blob_id=args_blob_id,
@@ -603,6 +621,14 @@ class FunctionCreationStatus:
                     f"Custom domain for {self.tag} => [magenta underline]{custom_domain.url}[/magenta underline]"
                 )
 
+        elif self.response.function.flash_service_urls:
+            for flash_service_url in self.response.function.flash_service_urls:
+                flash_service_url_status_row = self.resolver.add_status_row()
+                flash_service_url_status_row.finish(
+                    f"Created flash service endpoint for {self.tag} => "
+                    f"[magenta underline]{flash_service_url}[/magenta underline]"
+                )
+
         else:
             for warning in self.response.server_warnings:
                 self.status_row.warning(warning)