modal 1.0.3.dev10__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_serialization.py

@@ -6,14 +6,24 @@ import typing
 from inspect import Parameter
 from typing import Any
 
+from modal._traceback import extract_traceback
+from modal.config import config
+
+try:
+    import cbor2  # type: ignore
+except ImportError:  # pragma: no cover - optional dependency
+    cbor2 = None
+
+import google.protobuf.message
+
 from modal._utils.async_utils import synchronizer
 from modal_proto import api_pb2
 
 from ._object import _Object
 from ._type_manager import parameter_serde_registry, schema_registry
 from ._vendor import cloudpickle
-from .config import config, logger
-from .exception import DeserializationError, ExecutionError, InvalidError
+from .config import logger
+from .exception import DeserializationError, ExecutionError, InvalidError, SerializationError
 from .object import Object
 
 if typing.TYPE_CHECKING:
@@ -344,6 +354,12 @@ def _deserialize_asgi(asgi: api_pb2.Asgi) -> Any:
         return None
 
 
+def get_preferred_payload_format() -> "api_pb2.DataFormat.ValueType":
+    payload_format = (config.get("payload_format") or "pickle").lower()
+    data_format = api_pb2.DATA_FORMAT_CBOR if payload_format == "cbor" else api_pb2.DATA_FORMAT_PICKLE
+    return data_format
+
+
 def serialize_data_format(obj: Any, data_format: int) -> bytes:
     """Similar to serialize(), but supports other data formats."""
     if data_format == api_pb2.DATA_FORMAT_PICKLE:
@@ -353,6 +369,21 @@ def serialize_data_format(obj: Any, data_format: int) -> bytes:
     elif data_format == api_pb2.DATA_FORMAT_GENERATOR_DONE:
         assert isinstance(obj, api_pb2.GeneratorDone)
         return obj.SerializeToString(deterministic=True)
+    elif data_format == api_pb2.DATA_FORMAT_CBOR:
+        if cbor2 is None:
+            raise InvalidError("CBOR support requires the 'cbor2' package to be installed.")
+        try:
+            return cbor2.dumps(obj)
+        except cbor2.CBOREncodeTypeError:
+            try:
+                typename = f"{type(obj).__module__}.{type(obj).__name__}"
+            except Exception:
+                typename = str(type(obj))
+            raise SerializationError(
+                # TODO (elias): add documentation link for more information on this
+                f"Can not serialize type {typename} as cbor. If you need to use a custom data type, "
+                "try to serialize it yourself e.g. by using pickle.dumps(my_data)"
+            )
     else:
         raise InvalidError(f"Unknown data format {data_format!r}")
 
@@ -364,6 +395,10 @@ def deserialize_data_format(s: bytes, data_format: int, client) -> Any:
         return _deserialize_asgi(api_pb2.Asgi.FromString(s))
     elif data_format == api_pb2.DATA_FORMAT_GENERATOR_DONE:
         return api_pb2.GeneratorDone.FromString(s)
+    elif data_format == api_pb2.DATA_FORMAT_CBOR:
+        if cbor2 is None:
+            raise InvalidError("CBOR support requires the 'cbor2' package to be installed.")
+        return cbor2.loads(s)
     else:
         raise InvalidError(f"Unknown data format {data_format!r}")
 
@@ -470,10 +505,31 @@ def deserialize_params(serialized_params: bytes, function_def: api_pb2.Function,
         api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PICKLE,
     ):
         # legacy serialization format - pickle of `(args, kwargs)` w/ support for modal object arguments
-        param_args, param_kwargs = deserialize(serialized_params, _client)
+        try:
+            param_args, param_kwargs = deserialize(serialized_params, _client)
+        except DeserializationError as original_exc:
+            # Fallback in case of proto -> pickle downgrades of a parameter serialization format
+            # I.e. FunctionBindParams binding proto serialized params to a function defintion
+            # that now assumes pickled data according to class_parameter_info
+            param_args = ()
+            try:
+                param_kwargs = deserialize_proto_params(serialized_params)
+            except Exception:
+                raise original_exc
+
     elif function_def.class_parameter_info.format == api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PROTO:
         param_args = ()  # we use kwargs only for our implicit constructors
-        param_kwargs = deserialize_proto_params(serialized_params)
+        try:
+            param_kwargs = deserialize_proto_params(serialized_params)
+        except google.protobuf.message.DecodeError as original_exc:
+            # Fallback in case of pickle -> proto upgrades of a parameter serialization format
+            # I.e. FunctionBindParams binding pickle serialized params to a function defintion
+            # that now assumes proto data according to class_parameter_info
+            try:
+                param_args, param_kwargs = deserialize(serialized_params, _client)
+            except Exception:
+                raise original_exc
+
     else:
         raise ExecutionError(
             f"Unknown class parameter serialization format: {function_def.class_parameter_info.format}"
@@ -532,11 +588,16 @@ def get_callable_schema(
     callable: typing.Callable, *, is_web_endpoint: bool, ignore_first_argument: bool = False
 ) -> typing.Optional[api_pb2.FunctionSchema]:
     # ignore_first_argument can be used in case of unbound methods where we want to ignore the first (self) argument
-    if is_web_endpoint or not config.get("function_schemas"):
+    if is_web_endpoint:
         # we don't support schemas on web endpoints for now
         return None
 
-    sig = inspect.signature(callable)
+    try:
+        sig = inspect.signature(callable)
+    except Exception as e:
+        logger.debug(f"Error getting signature for function {callable}", exc_info=e)
+        return None
+
     # TODO: treat no return value annotation as None return?
     return_type_proto = schema_registry.get_proto_generic_type(sig.return_annotation)
     arguments = []
@@ -551,3 +612,26 @@ def get_callable_schema(
         arguments=arguments,
         return_type=return_type_proto,
     )
+
+
+def pickle_exception(exc: BaseException) -> bytes:
+    try:
+        return serialize(exc)
+    except Exception as serialization_exc:
+        # We can't always serialize exceptions.
+        err = f"Failed to serialize exception {exc} of type {type(exc)}: {serialization_exc}"
+        logger.info(err)
+        return serialize(SerializationError(err))
+
+
+def pickle_traceback(exc: BaseException, task_id: str) -> tuple[bytes, bytes]:
+    serialized_tb, tb_line_cache = b"", b""
+
+    try:
+        tb_dict, line_cache = extract_traceback(exc, task_id)
+        serialized_tb = serialize(tb_dict)
+        tb_line_cache = serialize(line_cache)
+    except Exception:
+        logger.info("Failed to serialize exception traceback.")
+
+    return serialized_tb, tb_line_cache

modal/_traceback.py

@@ -8,10 +8,12 @@ so that Rich is not a dependency of the container Client.
 import re
 import sys
 import traceback
+import typing
 import warnings
 from types import TracebackType
 from typing import Any, Iterable, Optional
 
+from modal.config import config, logger
 from modal_proto import api_pb2
 
 from ._vendor.tblib import Traceback as TBLibTraceback
@@ -115,7 +117,7 @@ def traceback_contains_remote_call(tb: Optional[TracebackType]) -> bool:
 def print_exception(exc: Optional[type[BaseException]], value: Optional[BaseException], tb: Optional[TracebackType]):
     """Add backwards compatibility for printing exceptions with "notes" for Python<3.11."""
     traceback.print_exception(exc, value, tb)
-    if sys.version_info < (3, 11) and value is not None:
+    if sys.version_info < (3, 11) and value is not None:  # type: ignore
         notes = getattr(value, "__notes__", [])
         print(*notes, sep="\n", file=sys.stderr)
 
@@ -127,3 +129,42 @@ def print_server_warnings(server_warnings: Iterable[api_pb2.Warning]):
     """
     for warning in server_warnings:
         warnings.warn_explicit(warning.message, ServerWarning, "<modal-server>", 0)
+
+
+# for some reason, the traceback cleanup here can't be moved into a context manager :(
+traceback_suppression_note = (
+    "Internal Modal traceback frames are suppressed for readability. Use MODAL_TRACEBACK=1 to show a full traceback."
+)
+
+
+class suppress_tb_frames:
+    def __init__(self, n: int):
+        self.n = n
+
+    def __enter__(self):
+        pass
+
+    def __exit__(
+        self, exc_type: Optional[type[BaseException]], exc: Optional[BaseException], tb: Optional[TracebackType]
+    ) -> typing.Literal[False]:
+        # *base* exceptions like CancelledError, SystemExit etc. can come from random places,
+        # so we don't suppress tracebacks for those
+        is_base_exception = not isinstance(exc, Exception)
+        if config.get("traceback") or exc_type is None or is_base_exception:
+            return False
+
+        # modify traceback on exception object
+        try:
+            final_tb = tb
+            for _ in range(self.n):
+                final_tb = final_tb.tb_next
+        except AttributeError:
+            logger.debug(f"Failed to suppress {self.n} traceback frames from {str(exc_type)} {str(exc)}")
+            raise
+
+        exc.with_traceback(final_tb)
+        notes = getattr(exc, "__notes__", [])
+        if traceback_suppression_note not in notes:
+            # .add_note was added in Python 3.11
+            notes.append(traceback_suppression_note)
+        return False

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