modal 1.1.5.dev66__py3-none-any.whl → 1.3.1.dev8__py3-none-any.whl

modal/_runtime/gpu_memory_snapshot.py

@@ -18,11 +18,14 @@ from modal.config import config, logger
 
 CUDA_CHECKPOINT_PATH: str = config.get("cuda_checkpoint_path")
 
-# Maximum total duration for an entire toggle operation.
-CUDA_CHECKPOINT_TOGGLE_TIMEOUT: float = 5 * 60.0
-
 # Maximum total duration for each individual `cuda-checkpoint` invocation.
-CUDA_CHECKPOINT_TIMEOUT: float = 90
+CUDA_CHECKPOINT_TIMEOUT: float = 3 * 60.0
+
+# Number of retries for each individual `cuda-checkpoint --toggle` invocation.
+CUDA_CHECKPOINT_TOGGLE_NUM_RETRIES: int = 3
+
+# Maximum total duration for an entire toggle operation.
+CUDA_CHECKPOINT_TOGGLE_TIMEOUT: float = CUDA_CHECKPOINT_TOGGLE_NUM_RETRIES * CUDA_CHECKPOINT_TIMEOUT
 
 
 class CudaCheckpointState(Enum):
@@ -58,7 +61,7 @@ class CudaCheckpointProcess:
 
         start_time = time.monotonic()
         retry_count = 0
-        max_retries = 3
+        max_retries = CUDA_CHECKPOINT_TOGGLE_NUM_RETRIES
 
         attempts = 0
         while self._should_continue_toggle(
@@ -201,8 +204,7 @@ class CudaCheckpointSession:
                 [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(pid)],
                 capture_output=True,
                 text=True,
-                # This should be quick since no checkpoint has taken place yet
-                timeout=5,
+                timeout=CUDA_CHECKPOINT_TIMEOUT,
             )
 
             # If the command succeeds (return code 0), this PID has a CUDA session

modal/_runtime/user_code_event_loop.py

@@ -0,0 +1,80 @@
+# Copyright Modal Labs 2022
+# ruff: noqa: E402
+import asyncio
+import signal
+import sys
+
+
+class UserCodeEventLoop:
+    """Run an async event loop as a context manager and handle signals.
+
+    This will run all *user supplied* async code, i.e. async functions, as well as async enter/exit managers
+
+    The following signals are handled while a coroutine is running on the event loop until
+    completion (and then handlers are deregistered):
+
+    - `SIGUSR1`: converted to an async task cancellation. Note that this only affects the event
+      loop, and the signal handler defined here doesn't run for sync functions.
+    - `SIGINT`: Unless the global signal handler has been set to SIGIGN, the loop's signal handler
+        is set to cancel the current task and raise KeyboardInterrupt to the caller.
+    """
+
+    def __enter__(self):
+        self.loop = asyncio.new_event_loop()
+        self.tasks = set()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.loop.run_until_complete(self.loop.shutdown_asyncgens())
+        if sys.version_info[:2] >= (3, 9):
+            self.loop.run_until_complete(self.loop.shutdown_default_executor())  # Introduced in Python 3.9
+
+        for task in self.tasks:
+            task.cancel()
+
+        self.loop.close()
+
+    def create_task(self, coro):
+        task = self.loop.create_task(coro)
+        self.tasks.add(task)
+        task.add_done_callback(self.tasks.discard)
+        return task
+
+    def run(self, coro):
+        task = asyncio.ensure_future(coro, loop=self.loop)
+        self._sigints = 0
+
+        def _sigint_handler():
+            # cancel the task in order to have run_until_complete return soon and
+            # prevent a bunch of unwanted tracebacks when shutting down the
+            # event loop.
+
+            # this basically replicates the sigint handler installed by asyncio.run()
+            self._sigints += 1
+            if self._sigints == 1:
+                # first sigint is graceful
+                task.cancel()
+                return
+
+            # this should normally not happen, but the second sigint would "hard kill" the event loop!
+            raise KeyboardInterrupt()
+
+        ignore_sigint = signal.getsignal(signal.SIGINT) == signal.SIG_IGN
+        if not ignore_sigint:
+            self.loop.add_signal_handler(signal.SIGINT, _sigint_handler)
+
+        # Before Python 3.9 there is no argument to Task.cancel
+        if sys.version_info[:2] >= (3, 9):
+            self.loop.add_signal_handler(signal.SIGUSR1, task.cancel, "Input was cancelled by user")
+        else:
+            self.loop.add_signal_handler(signal.SIGUSR1, task.cancel)
+
+        try:
+            return self.loop.run_until_complete(task)
+        except asyncio.CancelledError:
+            if self._sigints > 0:
+                raise KeyboardInterrupt()
+        finally:
+            self.loop.remove_signal_handler(signal.SIGUSR1)
+            if not ignore_sigint:
+                self.loop.remove_signal_handler(signal.SIGINT)

modal/_runtime/user_code_imports.py

@@ -1,20 +1,36 @@
 # Copyright Modal Labs 2024
 import importlib
+import inspect
+import os
+import signal
+import sys
 import typing
 from abc import ABCMeta, abstractmethod
+from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, Generator, Optional, Sequence
 
 import modal._object
 import modal._runtime.container_io_manager
 import modal.cls
 from modal import Function
 from modal._functions import _Function
+from modal._partial_function import (
+    _find_callables_for_obj,
+    _PartialFunctionFlags,
+)
+from modal._runtime.user_code_event_loop import UserCodeEventLoop
 from modal._utils.async_utils import synchronizer
-from modal._utils.function_utils import LocalFunctionError, is_async as get_is_async, is_global_object
+from modal._utils.function_utils import (
+    LocalFunctionError,
+    callable_has_non_self_params,
+    is_async as get_is_async,
+    is_global_object,
+)
 from modal.app import _App
 from modal.config import logger
 from modal.exception import ExecutionError, InvalidError
+from modal.experimental.flash import _FlashContainerEntry
 from modal_proto import api_pb2
 
 if typing.TYPE_CHECKING:
@@ -33,6 +49,68 @@ class FinalizedFunction:
     lifespan_manager: Optional["LifespanManager"] = None
 
 
+def call_lifecycle_functions(
+    event_loop: UserCodeEventLoop,
+    container_io_manager: Any,
+    funcs: Sequence[Callable[..., Any]],
+) -> None:
+    """Call function(s), can be sync or async, but any return values are ignored."""
+    with container_io_manager.handle_user_exception():
+        for func in funcs:
+            # We are deprecating parametrized exit methods but want to gracefully handle old code.
+            args = (None, None, None) if callable_has_non_self_params(func) else ()
+            res = func(*args)
+            if inspect.iscoroutine(res):
+                event_loop.run(res)
+
+
+@contextmanager
+def lifecycle_asgi(
+    event_loop: UserCodeEventLoop,
+    container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    finalized_functions: dict[str, FinalizedFunction],
+) -> Generator[None, None, None]:
+    lifespan_background_tasks = []
+    try:
+        for finalized_function in finalized_functions.values():
+            if finalized_function.lifespan_manager:
+                lifespan_background_tasks.append(
+                    event_loop.create_task(finalized_function.lifespan_manager.background_task())
+                )
+                with container_io_manager.handle_user_exception():
+                    event_loop.run(finalized_function.lifespan_manager.lifespan_startup())
+        yield
+    finally:
+        try:
+            # run lifespan shutdown for asgi apps
+            for finalized_function in finalized_functions.values():
+                if finalized_function.lifespan_manager:
+                    with container_io_manager.handle_user_exception():
+                        event_loop.run(finalized_function.lifespan_manager.lifespan_shutdown())
+        finally:
+            # no need to keep the lifespan asgi call around - we send it no more messages
+            for task in lifespan_background_tasks:
+                task.cancel()
+
+
+def disable_signals():
+    int_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+    usr1_handler = signal.signal(signal.SIGUSR1, signal.SIG_IGN)
+    return int_handler, usr1_handler
+
+
+def try_enable_signals(int_handler, usr1_handler):
+    if int_handler is not None and usr1_handler is not None:
+        signal.signal(signal.SIGINT, int_handler)
+        signal.signal(signal.SIGUSR1, usr1_handler)
+
+
+def volume_commit(
+    container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager", function_def: api_pb2.Function
+):
+    container_io_manager.volume_commit([v.volume_id for v in function_def.volume_mounts if v.allow_background_commits])
+
+
 class Service(metaclass=ABCMeta):
     """Common interface for singular functions and class-based "services"
 
@@ -44,12 +122,74 @@ class Service(metaclass=ABCMeta):
     user_cls_instance: Any
     app: "modal.app._App"
     service_deps: Optional[Sequence["modal._object._Object"]]
+    function_def: api_pb2.Function
 
     @abstractmethod
     def get_finalized_functions(
         self, fun_def: api_pb2.Function, container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager"
     ) -> dict[str, "FinalizedFunction"]: ...
 
+    @abstractmethod
+    @contextmanager
+    def lifecycle_presnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ) -> Generator[None, None, None]: ...
+
+    @abstractmethod
+    @contextmanager
+    def lifecycle_postsnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ) -> Generator[None, None, None]: ...
+
+    @contextmanager
+    def execution_context(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ) -> Generator[dict[str, "FinalizedFunction"], None, None]:
+        """
+        Manages the lifecycle of the user code:
+        1. Runs pre-snapshot 'enter' methods
+        2. Calls maybe_snapshot(container_io_manager, function_def)
+        3. Creates breakpoint wrapper
+        4. Runs post-snapshot 'enter' methods
+        5. Initializes finalized functions (and ASGI/WSGI lifespan)
+        6. Yield finalized_functions for execution
+        7. Handles cleanup (lifespan shutdown, 'exit' methods)
+        """
+        int_handler, usr1_handler = None, None
+        try:
+            # 1. Pre-snapshot Enter
+            with self.lifecycle_presnapshot(event_loop, container_io_manager):
+                # 2. Snapshot -- If this container is being used to create a checkpoint, checkpoint the container after
+                # global imports and initialization. Checkpointed containers run from this point onwards.
+                maybe_snapshot(container_io_manager, self.function_def)
+                # 3. Breakpoint wrapper
+                create_breakpoint_wrapper(container_io_manager)
+                # 4. Post-snapshot Enter
+                with self.lifecycle_postsnapshot(event_loop, container_io_manager):
+                    # Get Functions
+                    with container_io_manager.handle_user_exception():
+                        finalized_functions = self.get_finalized_functions(self.function_def, container_io_manager)
+                    # 5. Start ASGI lifespan
+                    with lifecycle_asgi(event_loop, container_io_manager, finalized_functions):
+                        # 6. Yield Finalized Functions
+                        try:
+                            yield finalized_functions
+                        finally:
+                            int_handler, usr1_handler = disable_signals()
+        finally:
+            # 9. Volume commit - runs OUTSIDE all lifecycle managers so exit handlers
+            # have a chance to write to disk before we commit volumes
+            try:
+                volume_commit(container_io_manager, self.function_def)
+            finally:
+                try_enable_signals(int_handler, usr1_handler)
+
 
 def construct_webhook_callable(
     user_defined_callable: Callable,
@@ -91,11 +231,36 @@ def construct_webhook_callable(
         raise InvalidError(f"Unrecognized web endpoint type {webhook_config.type}")
 
 
+def maybe_snapshot(
+    container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager", function_def: api_pb2.Function
+):
+    if function_def.is_checkpointing_function and os.environ.get("MODAL_ENABLE_SNAP_RESTORE") == "1":
+        container_io_manager.memory_snapshot()
+
+
+def create_breakpoint_wrapper(container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager"):
+    # Install hooks for interactive functions.
+    def breakpoint_wrapper():
+        # note: it would be nice to not have breakpoint_wrapper() included in the backtrace
+        container_io_manager.interact(from_breakpoint=True)
+        import pdb  # noqa: T100
+
+        current_frame = inspect.currentframe()
+        if current_frame is not None:
+            frame = current_frame.f_back
+            pdb.Pdb().set_trace(frame)
+        else:
+            raise RuntimeError("No current frame found")
+
+    sys.breakpointhook = breakpoint_wrapper
+
+
 @dataclass
 class ImportedFunction(Service):
     app: modal.app._App
     service_deps: Optional[Sequence["modal._object._Object"]]
     user_cls_instance = None
+    function_def: api_pb2.Function
 
     _user_defined_callable: Callable[..., Any]
 
@@ -138,6 +303,24 @@ class ImportedFunction(Service):
             )
         }
 
+    @contextmanager
+    def lifecycle_presnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ):
+        # This is a no-op for imported functions since @enter methods are not supported
+        yield
+
+    @contextmanager
+    def lifecycle_postsnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ):
+        # This is a no-op for imported functions since @enter methods are not supported
+        yield
+
 
 @dataclass
 class ImportedClass(Service):
@@ -146,6 +329,7 @@ class ImportedClass(Service):
     service_deps: Optional[Sequence["modal._object._Object"]]
 
     _partial_functions: dict[str, "modal._partial_function._PartialFunction"]
+    function_def: api_pb2.Function
 
     def get_finalized_functions(
         self, fun_def: api_pb2.Function, container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager"
@@ -189,6 +373,43 @@ class ImportedClass(Service):
             finalized_functions[method_name] = finalized_function
         return finalized_functions
 
+    @contextmanager
+    def lifecycle_presnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ):
+        # Identify all "enter" methods that need to run before we snapshot.
+        if not self.function_def.is_auto_snapshot:
+            pre_snapshot_methods = _find_callables_for_obj(
+                self.user_cls_instance, _PartialFunctionFlags.ENTER_PRE_SNAPSHOT
+            )
+            call_lifecycle_functions(event_loop, container_io_manager, list(pre_snapshot_methods.values()))
+        yield
+
+    @contextmanager
+    def lifecycle_postsnapshot(
+        self,
+        event_loop: UserCodeEventLoop,
+        container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    ):
+        flash_entry = _FlashContainerEntry(self.function_def.http_config)
+        # Identify the "enter" methods to run after resuming from a snapshot.
+        if not self.function_def.is_auto_snapshot:
+            post_snapshot_methods = _find_callables_for_obj(
+                self.user_cls_instance, _PartialFunctionFlags.ENTER_POST_SNAPSHOT
+            )
+            call_lifecycle_functions(event_loop, container_io_manager, list(post_snapshot_methods.values()))
+            flash_entry.enter()
+        try:
+            yield
+        finally:
+            if not self.function_def.is_auto_snapshot:
+                flash_entry.stop()
+                exit_methods = _find_callables_for_obj(self.user_cls_instance, _PartialFunctionFlags.EXIT)
+                call_lifecycle_functions(event_loop, container_io_manager, list(exit_methods.values()))
+                flash_entry.close()
+
 
 def get_user_class_instance(_cls: modal.cls._Cls, args: tuple[Any, ...], kwargs: dict[str, Any]) -> typing.Any:
     """Returns instance of the underlying class to be used as the `self`
@@ -244,7 +465,10 @@ def import_single_function_service(
     else:
         # Load the module dynamically
         module = importlib.import_module(function_def.module_name)
-        qual_name: str = function_def.function_name
+
+        # Fall back to function_name just to be safe around the migration
+        # Going forward, implementation_name should always be set
+        qual_name: str = function_def.implementation_name or function_def.function_name
 
         if not is_global_object(qual_name):
             raise LocalFunctionError("Attempted to load a function defined in a function scope")
@@ -266,9 +490,10 @@ def import_single_function_service(
             active_app = get_active_app_fallback(function_def)
 
     return ImportedFunction(
-        active_app,
-        service_deps,
-        user_defined_callable,
+        app=active_app,
+        service_deps=service_deps,
+        function_def=function_def,
+        _user_defined_callable=user_defined_callable,
     )
 
 
@@ -338,11 +563,12 @@ def import_class_service(
     user_cls_instance = get_user_class_instance(_cls, cls_args, cls_kwargs)
 
     return ImportedClass(
-        user_cls_instance,
-        active_app,
-        service_deps,
+        user_cls_instance=user_cls_instance,
+        app=active_app,
+        service_deps=service_deps,
         # TODO (elias/deven): instead of using method_partials here we should use a set of api_pb2.MethodDefinition
-        method_partials,
+        _partial_functions=method_partials,
+        function_def=function_def,
     )
 
 

modal/_serialization.py

@@ -420,7 +420,8 @@ def check_valid_cls_constructor_arg(key, obj):
     try:
         ClsConstructorPickler(buf).dump(obj)
         return True
-    except (AttributeError, ValueError):
+    except (AttributeError, ValueError, pickle.PicklingError):
+        # Python 3.14+ now raises an PicklingError for certain types of `dump` failures
         raise ValueError(
             f"Only pickle-able types are allowed in remote class constructors: argument {key} of type {type(obj)}."
         )

modal/_traceback.py

@@ -13,7 +13,7 @@ import warnings
 from types import TracebackType
 from typing import Any, Iterable, Optional
 
-from modal.config import config, logger
+from modal.config import config
 from modal_proto import api_pb2
 
 from ._vendor.tblib import Traceback as TBLibTraceback
@@ -119,7 +119,7 @@ def print_exception(exc: Optional[type[BaseException]], value: Optional[BaseExce
     traceback.print_exception(exc, value, tb)
     if sys.version_info < (3, 11) and value is not None:  # type: ignore
         notes = getattr(value, "__notes__", [])
-        print(*notes, sep="\n", file=sys.stderr)
+        print(*notes, sep="\n", file=sys.stderr)  # noqa: T201
 
 
 def print_server_warnings(server_warnings: Iterable[api_pb2.Warning]):
@@ -137,10 +137,7 @@ traceback_suppression_note = (
 )
 
 
-class suppress_tb_frames:
-    def __init__(self, n: int):
-        self.n = n
-
+class suppress_tb_frame:
     def __enter__(self):
         pass
 
@@ -154,13 +151,7 @@ class suppress_tb_frames:
             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
+        final_tb = tb.tb_next if tb is not None else tb
 
         exc.with_traceback(final_tb)
         notes = getattr(exc, "__notes__", [])

modal/_tunnel.py

@@ -5,14 +5,13 @@ from collections.abc import AsyncIterator
 from dataclasses import dataclass
 from typing import Optional
 
-from grpclib import GRPCError, Status
 from synchronicity.async_wrap import asynccontextmanager
 
 from modal_proto import api_pb2
 
 from ._utils.async_utils import synchronize_api
 from .client import _Client
-from .exception import InvalidError, RemoteError
+from .exception import AlreadyExistsError, InvalidError, RemoteError, ServiceError
 
 
 @dataclass(frozen=True)
@@ -51,13 +50,17 @@ class Tunnel:
 
 
 @asynccontextmanager
-async def _forward(port: int, *, unencrypted: bool = False, client: Optional[_Client] = None) -> AsyncIterator[Tunnel]:
+async def _forward(
+    port: int, *, unencrypted: bool = False, h2_enabled: bool = False, client: Optional[_Client] = None
+) -> AsyncIterator[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.
 
+    If `h2_enabled` is set, the TLS server will advertise support for HTTP/2.
+
     **Important:** This is an experimental API which may change in the future.
 
     **Usage:**
@@ -168,6 +171,8 @@ async def _forward(port: int, *, unencrypted: bool = False, client: Optional[_Cl
         raise InvalidError(f"The port argument should be an int, not {port!r}")
     if port < 1 or port > 65535:
         raise InvalidError(f"Invalid port number {port}")
+    if h2_enabled and unencrypted:
+        raise InvalidError("H2 can only be used with encrypted ports")
 
     if not client:
         client = await _Client.from_env()
@@ -175,15 +180,15 @@ async def _forward(port: int, *, unencrypted: bool = False, client: Optional[_Cl
     if client.client_type != api_pb2.CLIENT_TYPE_CONTAINER:
         raise InvalidError("Forwarding ports only works inside a Modal container")
 
+    tunnel_type = api_pb2.TUNNEL_TYPE_H2 if h2_enabled else api_pb2.TUNNEL_TYPE_UNSPECIFIED
     try:
-        response = await client.stub.TunnelStart(api_pb2.TunnelStartRequest(port=port, unencrypted=unencrypted))
-    except GRPCError as exc:
-        if exc.status == Status.ALREADY_EXISTS:
-            raise InvalidError(f"Port {port} is already forwarded")
-        elif exc.status == Status.UNAVAILABLE:
-            raise RemoteError("Relay server is unavailable") from exc
-        else:
-            raise
+        response = await client.stub.TunnelStart(
+            api_pb2.TunnelStartRequest(port=port, unencrypted=unencrypted, tunnel_type=tunnel_type)
+        )
+    except AlreadyExistsError as exc:
+        raise InvalidError(f"Port {port} is already forwarded")
+    except ServiceError as exc:
+        raise RemoteError("Relay server is unavailable") from exc
 
     try:
         yield Tunnel(response.host, response.port, response.unencrypted_host, response.unencrypted_port)

modal/_tunnel.pyi

@@ -54,7 +54,11 @@ class Tunnel:
         ...
 
 def _forward(
-    port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client._Client] = None
+    port: int,
+    *,
+    unencrypted: bool = False,
+    h2_enabled: bool = False,
+    client: typing.Optional[modal.client._Client] = None,
 ) -> typing.AsyncContextManager[Tunnel]:
     '''Expose a port publicly from inside a running Modal container, with TLS.
 
@@ -62,6 +66,8 @@ def _forward(
     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.
 
+    If `h2_enabled` is set, the TLS server will advertise support for HTTP/2.
+
     **Important:** This is an experimental API which may change in the future.
 
     **Usage:**
@@ -171,7 +177,13 @@ def _forward(
 
 class __forward_spec(typing_extensions.Protocol):
     def __call__(
-        self, /, port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client.Client] = None
+        self,
+        /,
+        port: int,
+        *,
+        unencrypted: bool = False,
+        h2_enabled: bool = False,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[Tunnel]:
         '''Expose a port publicly from inside a running Modal container, with TLS.
 
@@ -179,6 +191,8 @@ class __forward_spec(typing_extensions.Protocol):
         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.
 
+        If `h2_enabled` is set, the TLS server will advertise support for HTTP/2.
+
         **Important:** This is an experimental API which may change in the future.
 
         **Usage:**
@@ -287,7 +301,13 @@ class __forward_spec(typing_extensions.Protocol):
         ...
 
     def aio(
-        self, /, port: int, *, unencrypted: bool = False, client: typing.Optional[modal.client.Client] = None
+        self,
+        /,
+        port: int,
+        *,
+        unencrypted: bool = False,
+        h2_enabled: bool = False,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> typing.AsyncContextManager[Tunnel]:
         '''Expose a port publicly from inside a running Modal container, with TLS.
 
@@ -295,6 +315,8 @@ class __forward_spec(typing_extensions.Protocol):
         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.
 
+        If `h2_enabled` is set, the TLS server will advertise support for HTTP/2.
+
         **Important:** This is an experimental API which may change in the future.
 
         **Usage:**