modal 0.62.115__py3-none-any.whl → 0.72.13__py3-none-any.whl

modal/partial_function.py

@@ -1,65 +1,114 @@
 # Copyright Modal Labs 2023
 import enum
+import inspect
+import typing
+from collections.abc import Coroutine, Iterable
 from typing import (
     Any,
     Callable,
-    Dict,
-    Iterable,
-    List,
     Optional,
-    Type,
     Union,
 )
 
+import typing_extensions
+
 from modal_proto import api_pb2
 
 from ._utils.async_utils import synchronize_api, synchronizer
-from ._utils.function_utils import method_has_params
+from ._utils.deprecation import deprecation_error, deprecation_warning
+from ._utils.function_utils import callable_has_non_self_non_default_params, callable_has_non_self_params
 from .config import logger
-from .exception import InvalidError, deprecation_warning
+from .exception import InvalidError
 from .functions import _Function
 
+MAX_MAX_BATCH_SIZE = 1000
+MAX_BATCH_WAIT_MS = 10 * 60 * 1000  # 10 minutes
+
 
 class _PartialFunctionFlags(enum.IntFlag):
     FUNCTION: int = 1
     BUILD: int = 2
-    ENTER_PRE_CHECKPOINT: int = 4
-    ENTER_POST_CHECKPOINT: int = 8
+    ENTER_PRE_SNAPSHOT: int = 4
+    ENTER_POST_SNAPSHOT: int = 8
     EXIT: int = 16
+    BATCHED: int = 32
+    CLUSTERED: int = 64  # Experimental: Clustered functions
+
+    @staticmethod
+    def all() -> int:
+        return ~_PartialFunctionFlags(0)
+
+
+P = typing_extensions.ParamSpec("P")
+ReturnType = typing_extensions.TypeVar("ReturnType", covariant=True)
+OriginalReturnType = typing_extensions.TypeVar("OriginalReturnType", covariant=True)
 
 
-class _PartialFunction:
-    """Intermediate function, produced by @method or @web_endpoint"""
+class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
+    """Intermediate function, produced by @enter, @build, @method, @web_endpoint, or @batched"""
 
-    raw_f: Callable[..., Any]
+    raw_f: Callable[P, ReturnType]
     flags: _PartialFunctionFlags
     webhook_config: Optional[api_pb2.WebhookConfig]
-    is_generator: Optional[bool]
+    is_generator: bool
     keep_warm: Optional[int]
+    batch_max_size: Optional[int]
+    batch_wait_ms: Optional[int]
+    force_build: bool
+    cluster_size: Optional[int]  # Experimental: Clustered functions
+    build_timeout: Optional[int]
 
     def __init__(
         self,
-        raw_f: Callable[..., Any],
+        raw_f: Callable[P, ReturnType],
         flags: _PartialFunctionFlags,
         webhook_config: Optional[api_pb2.WebhookConfig] = None,
         is_generator: Optional[bool] = None,
         keep_warm: Optional[int] = None,
+        batch_max_size: Optional[int] = None,
+        batch_wait_ms: Optional[int] = None,
+        cluster_size: Optional[int] = None,  # Experimental: Clustered functions
+        force_build: bool = False,
+        build_timeout: Optional[int] = None,
     ):
         self.raw_f = raw_f
         self.flags = flags
         self.webhook_config = webhook_config
-        self.is_generator = is_generator
+        if is_generator is None:
+            # auto detect - doesn't work if the function *returns* a generator
+            final_is_generator = inspect.isgeneratorfunction(raw_f) or inspect.isasyncgenfunction(raw_f)
+        else:
+            final_is_generator = is_generator
+
+        self.is_generator = final_is_generator
         self.keep_warm = keep_warm
         self.wrapped = False  # Make sure that this was converted into a FunctionHandle
+        self.batch_max_size = batch_max_size
+        self.batch_wait_ms = batch_wait_ms
+        self.cluster_size = cluster_size  # Experimental: Clustered functions
+        self.force_build = force_build
+        self.build_timeout = build_timeout
 
-    def __get__(self, obj, objtype=None) -> _Function:
-        # This only happens inside user methods when they refer to other methods
+    def __get__(self, obj, objtype=None) -> _Function[P, ReturnType, OriginalReturnType]:
         k = self.raw_f.__name__
-        if obj:  # Cls().fun
-            function = getattr(obj, "_modal_functions")[k]
-        else:  # Cls.fun
-            function = getattr(objtype, "_modal_functions")[k]
-        return function
+        if obj:  # accessing the method on an instance of a class, e.g. `MyClass().fun``
+            if hasattr(obj, "_modal_functions"):
+                # This happens inside "local" user methods when they refer to other methods,
+                # e.g. Foo().parent_method.remote() calling self.other_method.remote()
+                return getattr(obj, "_modal_functions")[k]
+            else:
+                # special edge case: referencing a method of an instance of an
+                # unwrapped class (not using app.cls()) with @methods
+                # not sure what would be useful here, but let's return a bound version of the underlying function,
+                # since the class is just a vanilla class at this point
+                # This wouldn't let the user access `.remote()` and `.local()` etc. on the function
+                return self.raw_f.__get__(obj, objtype)
+
+        else:  # accessing a method directly on the class, e.g. `MyClass.fun`
+            # This happens mainly during serialization of the wrapped underlying class of a Cls
+            # since we don't have the instance info here we just return the PartialFunction itself
+            # to let it be bound to a variable and become a Function later on
+            return self  # type: ignore  # this returns a PartialFunction in a special internal case
 
     def __del__(self):
         if (self.flags & _PartialFunctionFlags.FUNCTION) and self.wrapped is False:
@@ -76,76 +125,70 @@ class _PartialFunction:
             flags=(self.flags | flags),
             webhook_config=self.webhook_config,
             keep_warm=self.keep_warm,
+            batch_max_size=self.batch_max_size,
+            batch_wait_ms=self.batch_wait_ms,
+            force_build=self.force_build,
+            build_timeout=self.build_timeout,
         )
 
 
 PartialFunction = synchronize_api(_PartialFunction)
 
 
-def _find_partial_methods_for_cls(user_cls: Type, flags: _PartialFunctionFlags) -> Dict[str, _PartialFunction]:
-    """Grabs all method on a user class"""
-    partial_functions: Dict[str, PartialFunction] = {}
-    for parent_cls in user_cls.mro():
+def _find_partial_methods_for_user_cls(user_cls: type[Any], flags: int) -> dict[str, _PartialFunction]:
+    """Grabs all method on a user class, and returns partials. Includes legacy methods."""
+
+    partial_functions: dict[str, _PartialFunction] = {}
+    for parent_cls in reversed(user_cls.mro()):
         if parent_cls is not object:
             for k, v in parent_cls.__dict__.items():
-                if isinstance(v, PartialFunction):
-                    partial_function = synchronizer._translate_in(v)  # TODO: remove need for?
-                    if partial_function.flags & flags:
-                        partial_functions[k] = partial_function
+                if isinstance(v, PartialFunction):  # type: ignore[reportArgumentType]   # synchronicity wrapper types
+                    _partial_function: _PartialFunction = typing.cast(_PartialFunction, synchronizer._translate_in(v))
+                    if _partial_function.flags & flags:
+                        partial_functions[k] = _partial_function
 
     return partial_functions
 
 
-def _find_callables_for_cls(user_cls: Type, flags: _PartialFunctionFlags) -> Dict[str, Callable]:
-    """Grabs all method on a user class, and returns callables. Includes legacy methods."""
-    functions: Dict[str, Callable] = {}
-
-    # Build up a list of legacy attributes to check
-    check_attrs: List[str] = []
-    if flags & _PartialFunctionFlags.BUILD:
-        check_attrs += ["__build__", "__abuild__"]
-    if flags & _PartialFunctionFlags.ENTER_POST_CHECKPOINT:
-        check_attrs += ["__enter__", "__aenter__"]
-    if flags & _PartialFunctionFlags.EXIT:
-        check_attrs += ["__exit__", "__aexit__"]
-
-    # Grab legacy lifecycle methods
-    for attr in check_attrs:
-        if hasattr(user_cls, attr):
-            suggested = attr.strip("_")
-            if is_async := suggested.startswith("a"):
-                suggested = suggested[1:]
-            async_suggestion = " (on an async method)" if is_async else ""
-            message = (
-                f"Using `{attr}` methods for class lifecycle management is deprecated."
-                f" Please try using the `modal.{suggested}` decorator{async_suggestion} instead."
-                " See https://modal.com/docs/guide/lifecycle-functions for more information."
-            )
-            deprecation_warning((2024, 2, 21), message, show_source=True)
-            functions[attr] = getattr(user_cls, attr)
+def _find_callables_for_obj(user_obj: Any, flags: int) -> dict[str, Callable[..., Any]]:
+    """Grabs all methods for an object, and binds them to the class"""
+    user_cls: type = type(user_obj)
+    return {k: pf.raw_f.__get__(user_obj) for k, pf in _find_partial_methods_for_user_cls(user_cls, flags).items()}
 
-    # Grab new decorator-based methods
-    for k, pf in _find_partial_methods_for_cls(user_cls, flags).items():
-        functions[k] = pf.raw_f
 
-    return functions
+class _MethodDecoratorType:
+    @typing.overload
+    def __call__(
+        self, func: PartialFunction[typing_extensions.Concatenate[Any, P], ReturnType, OriginalReturnType]
+    ) -> PartialFunction[P, ReturnType, OriginalReturnType]:
+        ...
 
+    @typing.overload
+    def __call__(
+        self, func: Callable[typing_extensions.Concatenate[Any, P], Coroutine[Any, Any, ReturnType]]
+    ) -> PartialFunction[P, ReturnType, Coroutine[Any, Any, ReturnType]]:
+        ...
 
-def _find_callables_for_obj(user_obj: Any, flags: _PartialFunctionFlags) -> Dict[str, Callable]:
-    """Grabs all methods for an object, and binds them to the class"""
-    user_cls: Type = type(user_obj)
-    return {k: meth.__get__(user_obj) for k, meth in _find_callables_for_cls(user_cls, flags).items()}
+    @typing.overload
+    def __call__(
+        self, func: Callable[typing_extensions.Concatenate[Any, P], ReturnType]
+    ) -> PartialFunction[P, ReturnType, ReturnType]:
+        ...
+
+    def __call__(self, func):
+        ...
 
 
+# TODO(elias): fix support for coroutine type unwrapping for methods (static typing)
 def _method(
     _warn_parentheses_missing=None,
     *,
     # Set this to True if it's a non-generator function returning
     # a [sync/async] generator object
     is_generator: Optional[bool] = None,
-    keep_warm: Optional[int] = None,  # An optional number of containers to always keep warm.
-) -> Callable[[Callable[..., Any]], _PartialFunction]:
-    """Decorator for methods that should be transformed into a Modal Function registered against this class's app.
+    keep_warm: Optional[int] = None,  # Deprecated: Use keep_warm on @app.cls() instead
+) -> _MethodDecoratorType:
+    """Decorator for methods that should be transformed into a Modal Function registered against this class's App.
 
     **Usage:**
 
@@ -158,23 +201,42 @@ def _method(
             ...
     ```
     """
-    if _warn_parentheses_missing:
+    if _warn_parentheses_missing is not None:
         raise InvalidError("Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@method()`.")
 
+    if keep_warm is not None:
+        deprecation_warning(
+            (2024, 6, 10),
+            (
+                "`keep_warm=` is no longer supported per-method on Modal classes. "
+                "All methods and web endpoints of a class use the same set of containers now. "
+                "Use keep_warm via the @app.cls() decorator instead. "
+            ),
+            pending=True,
+        )
+
     def wrapper(raw_f: Callable[..., Any]) -> _PartialFunction:
+        nonlocal is_generator
         if isinstance(raw_f, _PartialFunction) and raw_f.webhook_config:
             raw_f.wrapped = True  # suppress later warning
             raise InvalidError(
-                "Web endpoints on classes should not be wrapped by `@method`. Suggestion: remove the `@method` decorator."
+                "Web endpoints on classes should not be wrapped by `@method`. "
+                "Suggestion: remove the `@method` decorator."
+            )
+        if isinstance(raw_f, _PartialFunction) and raw_f.batch_max_size is not None:
+            raw_f.wrapped = True  # suppress later warning
+            raise InvalidError(
+                "Batched function on classes should not be wrapped by `@method`. "
+                "Suggestion: remove the `@method` decorator."
             )
         return _PartialFunction(raw_f, _PartialFunctionFlags.FUNCTION, is_generator=is_generator, keep_warm=keep_warm)
 
     return wrapper
 
 
-def _parse_custom_domains(custom_domains: Optional[Iterable[str]] = None) -> List[api_pb2.CustomDomainConfig]:
+def _parse_custom_domains(custom_domains: Optional[Iterable[str]] = None) -> list[api_pb2.CustomDomainConfig]:
     assert not isinstance(custom_domains, str), "custom_domains must be `Iterable[str]` but is `str` instead."
-    _custom_domains: List[api_pb2.CustomDomainConfig] = []
+    _custom_domains: list[api_pb2.CustomDomainConfig] = []
     if custom_domains is not None:
         for custom_domain in custom_domains:
             _custom_domains.append(api_pb2.CustomDomainConfig(name=custom_domain))
@@ -187,11 +249,13 @@ def _web_endpoint(
     *,
     method: str = "GET",  # REST method for the created endpoint.
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
-    wait_for_response: bool = True,  # Whether requests should wait for and return the function response.
+    docs: bool = False,  # Whether to enable interactive documentation for this endpoint at /docs.
     custom_domains: Optional[
         Iterable[str]
     ] = None,  # Create an endpoint using a custom domain fully-qualified domain name (FQDN).
-) -> Callable[[Callable[..., Any]], _PartialFunction]:
+    requires_proxy_auth: bool = False,  # Require Proxy-Authorization HTTP Headers on requests
+    wait_for_response: bool = True,  # DEPRECATED: this must always be True now
+) -> Callable[[Callable[P, ReturnType]], _PartialFunction[P, ReturnType, ReturnType]]:
     """Register a basic web endpoint with this application.
 
     This is the simple way to create a web endpoint on Modal. The function
@@ -209,7 +273,7 @@ def _web_endpoint(
     if isinstance(_warn_parentheses_missing, str):
         # Probably passing the method string as a positional argument.
         raise InvalidError('Positional arguments are not allowed. Suggestion: `@web_endpoint(method="GET")`.')
-    elif _warn_parentheses_missing:
+    elif _warn_parentheses_missing is not None:
         raise InvalidError(
             "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@web_endpoint()`."
         )
@@ -222,9 +286,11 @@ def _web_endpoint(
                 "@app.function()\n@app.web_endpoint()\ndef my_webhook():\n    ..."
             )
         if not wait_for_response:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_TRIGGER
-        else:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_AUTO  # the default
+            deprecation_error(
+                (2024, 5, 13),
+                "wait_for_response=False has been deprecated on web endpoints. See "
+                "https://modal.com/docs/guide/webhook-timeouts#polling-solutions for alternatives.",
+            )
 
         # self._loose_webhook_configs.add(raw_f)
 
@@ -234,9 +300,11 @@ def _web_endpoint(
             api_pb2.WebhookConfig(
                 type=api_pb2.WEBHOOK_TYPE_FUNCTION,
                 method=method,
+                web_endpoint_docs=docs,
                 requested_suffix=label,
-                async_mode=_response_mode,
+                async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
                 custom_domains=_parse_custom_domains(custom_domains),
+                requires_proxy_auth=requires_proxy_auth,
             ),
         )
 
@@ -247,8 +315,9 @@ def _asgi_app(
     _warn_parentheses_missing=None,
     *,
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
-    wait_for_response: bool = True,  # Whether requests should wait for and return the function response.
     custom_domains: Optional[Iterable[str]] = None,  # Deploy this endpoint on a custom domain.
+    requires_proxy_auth: bool = False,  # Require Proxy-Authorization HTTP Headers on requests
+    wait_for_response: bool = True,  # DEPRECATED: this must always be True now
 ) -> Callable[[Callable[..., Any]], _PartialFunction]:
     """Decorator for registering an ASGI app with a Modal function.
 
@@ -273,16 +342,35 @@ def _asgi_app(
     """
     if isinstance(_warn_parentheses_missing, str):
         raise InvalidError('Positional arguments are not allowed. Suggestion: `@asgi_app(label="foo")`.')
-    elif _warn_parentheses_missing:
+    elif _warn_parentheses_missing is not None:
         raise InvalidError(
             "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@asgi_app()`."
         )
 
     def wrapper(raw_f: Callable[..., Any]) -> _PartialFunction:
+        if callable_has_non_self_params(raw_f):
+            if callable_has_non_self_non_default_params(raw_f):
+                raise InvalidError(
+                    f"ASGI app function {raw_f.__name__} can't have parameters. See https://modal.com/docs/guide/webhooks#asgi."
+                )
+            else:
+                deprecation_warning(
+                    (2024, 9, 4),
+                    f"ASGI app function {raw_f.__name__} has default parameters, but shouldn't have any parameters - "
+                    f"Modal will drop support for default parameters in a future release.",
+                )
+
+        if inspect.iscoroutinefunction(raw_f):
+            raise InvalidError(
+                f"ASGI app function {raw_f.__name__} is an async function. Only sync Python functions are supported."
+            )
+
         if not wait_for_response:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_TRIGGER
-        else:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_AUTO  # the default
+            deprecation_error(
+                (2024, 5, 13),
+                "wait_for_response=False has been deprecated on web endpoints. See "
+                "https://modal.com/docs/guide/webhook-timeouts#polling-solutions for alternatives",
+            )
 
         return _PartialFunction(
             raw_f,
@@ -290,8 +378,9 @@ def _asgi_app(
             api_pb2.WebhookConfig(
                 type=api_pb2.WEBHOOK_TYPE_ASGI_APP,
                 requested_suffix=label,
-                async_mode=_response_mode,
+                async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
                 custom_domains=_parse_custom_domains(custom_domains),
+                requires_proxy_auth=requires_proxy_auth,
             ),
         )
 
@@ -302,14 +391,16 @@ def _wsgi_app(
     _warn_parentheses_missing=None,
     *,
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
-    wait_for_response: bool = True,  # Whether requests should wait for and return the function response.
     custom_domains: Optional[Iterable[str]] = None,  # Deploy this endpoint on a custom domain.
+    requires_proxy_auth: bool = False,  # Require Proxy-Authorization HTTP Headers on requests
+    wait_for_response: bool = True,  # DEPRECATED: this must always be True now
 ) -> Callable[[Callable[..., Any]], _PartialFunction]:
     """Decorator for registering a WSGI app with a Modal function.
 
     Web Server Gateway Interface (WSGI) is a standard for synchronous Python web apps.
-    It has been [succeeded by the ASGI interface](https://asgi.readthedocs.io/en/latest/introduction.html#wsgi-compatibility) which is compatible with ASGI and supports
-    additional functionality such as web sockets. Modal supports ASGI via [`asgi_app`](/docs/reference/modal.asgi_app).
+    It has been [succeeded by the ASGI interface](https://asgi.readthedocs.io/en/latest/introduction.html#wsgi-compatibility)
+    which is compatible with ASGI and supports additional functionality such as web sockets.
+    Modal supports ASGI via [`asgi_app`](/docs/reference/modal.asgi_app).
 
     **Usage:**
 
@@ -327,16 +418,35 @@ def _wsgi_app(
     """
     if isinstance(_warn_parentheses_missing, str):
         raise InvalidError('Positional arguments are not allowed. Suggestion: `@wsgi_app(label="foo")`.')
-    elif _warn_parentheses_missing:
+    elif _warn_parentheses_missing is not None:
         raise InvalidError(
             "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@wsgi_app()`."
         )
 
     def wrapper(raw_f: Callable[..., Any]) -> _PartialFunction:
+        if callable_has_non_self_params(raw_f):
+            if callable_has_non_self_non_default_params(raw_f):
+                raise InvalidError(
+                    f"WSGI app function {raw_f.__name__} can't have parameters. See https://modal.com/docs/guide/webhooks#wsgi."
+                )
+            else:
+                deprecation_warning(
+                    (2024, 9, 4),
+                    f"WSGI app function {raw_f.__name__} has default parameters, but shouldn't have any parameters - "
+                    f"Modal will drop support for default parameters in a future release.",
+                )
+
+        if inspect.iscoroutinefunction(raw_f):
+            raise InvalidError(
+                f"WSGI app function {raw_f.__name__} is an async function. Only sync Python functions are supported."
+            )
+
         if not wait_for_response:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_TRIGGER
-        else:
-            _response_mode = api_pb2.WEBHOOK_ASYNC_MODE_AUTO  # the default
+            deprecation_error(
+                (2024, 5, 13),
+                "wait_for_response=False has been deprecated on web endpoints. See "
+                "https://modal.com/docs/guide/webhook-timeouts#polling-solutions for alternatives",
+            )
 
         return _PartialFunction(
             raw_f,
@@ -344,8 +454,9 @@ def _wsgi_app(
             api_pb2.WebhookConfig(
                 type=api_pb2.WEBHOOK_TYPE_WSGI_APP,
                 requested_suffix=label,
-                async_mode=_response_mode,
+                async_mode=api_pb2.WEBHOOK_ASYNC_MODE_AUTO,
                 custom_domains=_parse_custom_domains(custom_domains),
+                requires_proxy_auth=requires_proxy_auth,
             ),
         )
 
@@ -358,6 +469,7 @@ def _web_server(
     startup_timeout: float = 5.0,  # Maximum number of seconds to wait for the web server to start.
     label: Optional[str] = None,  # Label for created endpoint. Final subdomain will be <workspace>--<label>.modal.run.
     custom_domains: Optional[Iterable[str]] = None,  # Deploy this endpoint on a custom domain.
+    requires_proxy_auth: bool = False,  # Require Proxy-Authorization HTTP Headers on requests
 ) -> Callable[[Callable[..., Any]], _PartialFunction]:
     """Decorator that registers an HTTP web server inside the container.
 
@@ -401,6 +513,7 @@ def _web_server(
                 custom_domains=_parse_custom_domains(custom_domains),
                 web_server_port=port,
                 web_server_startup_timeout=startup_timeout,
+                requires_proxy_auth=requires_proxy_auth,
             ),
         )
 
@@ -414,7 +527,7 @@ def _disallow_wrapping_method(f: _PartialFunction, wrapper: str) -> None:
 
 
 def _build(
-    _warn_parentheses_missing=None,
+    _warn_parentheses_missing=None, *, force: bool = False, timeout: int = 86400
 ) -> Callable[[Union[Callable[[Any], Any], _PartialFunction]], _PartialFunction]:
     """
     Decorator for methods that should execute at _build time_ to create a new layer
@@ -436,15 +549,17 @@ def _build(
             LlamaTokenizer.from_pretrained(base_model)
     ```
     """
-    if _warn_parentheses_missing:
+    if _warn_parentheses_missing is not None:
         raise InvalidError("Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@build()`.")
 
     def wrapper(f: Union[Callable[[Any], Any], _PartialFunction]) -> _PartialFunction:
         if isinstance(f, _PartialFunction):
             _disallow_wrapping_method(f, "build")
+            f.force_build = force
+            f.build_timeout = timeout
             return f.add_flags(_PartialFunctionFlags.BUILD)
         else:
-            return _PartialFunction(f, _PartialFunctionFlags.BUILD)
+            return _PartialFunction(f, _PartialFunctionFlags.BUILD, force_build=force, build_timeout=timeout)
 
     return wrapper
 
@@ -457,13 +572,13 @@ def _enter(
     """Decorator for methods which should be executed when a new container is started.
 
     See the [lifeycle function guide](https://modal.com/docs/guide/lifecycle-functions#enter) for more information."""
-    if _warn_parentheses_missing:
+    if _warn_parentheses_missing is not None:
         raise InvalidError("Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@enter()`.")
 
     if snap:
-        flag = _PartialFunctionFlags.ENTER_PRE_CHECKPOINT
+        flag = _PartialFunctionFlags.ENTER_PRE_SNAPSHOT
     else:
-        flag = _PartialFunctionFlags.ENTER_POST_CHECKPOINT
+        flag = _PartialFunctionFlags.ENTER_POST_SNAPSHOT
 
     def wrapper(f: Union[Callable[[Any], Any], _PartialFunction]) -> _PartialFunction:
         if isinstance(f, _PartialFunction):
@@ -476,10 +591,12 @@ def _enter(
 
 
 ExitHandlerType = Union[
+    # NOTE: return types of these callables should be `Union[None, Awaitable[None]]` but
+    #       synchronicity type stubs would strip Awaitable so we use Any for now
     # Original, __exit__ style method signature (now deprecated)
-    Callable[[Any, Optional[Type[BaseException]], Optional[BaseException], Any], None],
+    Callable[[Any, Optional[type[BaseException]], Optional[BaseException], Any], Any],
     # Forward-looking unparameterized method
-    Callable[[Any], None],
+    Callable[[Any], Any],
 ]
 
 
@@ -487,23 +604,70 @@ def _exit(_warn_parentheses_missing=None) -> Callable[[ExitHandlerType], _Partia
     """Decorator for methods which should be executed when a container is about to exit.
 
     See the [lifeycle function guide](https://modal.com/docs/guide/lifecycle-functions#exit) for more information."""
-    if _warn_parentheses_missing:
+    if _warn_parentheses_missing is not None:
         raise InvalidError("Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@exit()`.")
 
     def wrapper(f: ExitHandlerType) -> _PartialFunction:
         if isinstance(f, _PartialFunction):
             _disallow_wrapping_method(f, "exit")
-        if method_has_params(f):
-            message = (
-                "Support for decorating parameterized methods with `@exit` has been deprecated."
-                " To avoid future errors, please update your code by removing the parameters."
-            )
-            deprecation_warning((2024, 2, 23), message)
+
         return _PartialFunction(f, _PartialFunctionFlags.EXIT)
 
     return wrapper
 
 
+def _batched(
+    _warn_parentheses_missing=None,
+    *,
+    max_batch_size: int,
+    wait_ms: int,
+) -> Callable[[Callable[..., Any]], _PartialFunction]:
+    """Decorator for functions or class methods that should be batched.
+
+    **Usage**
+
+    ```python notest
+    @app.function()
+    @modal.batched(max_batch_size=4, wait_ms=1000)
+    async def batched_multiply(xs: list[int], ys: list[int]) -> list[int]:
+        return [x * y for x, y in zip(xs, xs)]
+
+    # call batched_multiply with individual inputs
+    batched_multiply.remote.aio(2, 100)
+    ```
+
+    See the [dynamic batching guide](https://modal.com/docs/guide/dynamic-batching) for more information.
+    """
+    if _warn_parentheses_missing is not None:
+        raise InvalidError(
+            "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@batched()`."
+        )
+    if max_batch_size < 1:
+        raise InvalidError("max_batch_size must be a positive integer.")
+    if max_batch_size >= MAX_MAX_BATCH_SIZE:
+        raise InvalidError(f"max_batch_size must be less than {MAX_MAX_BATCH_SIZE}.")
+    if wait_ms < 0:
+        raise InvalidError("wait_ms must be a non-negative integer.")
+    if wait_ms >= MAX_BATCH_WAIT_MS:
+        raise InvalidError(f"wait_ms must be less than {MAX_BATCH_WAIT_MS}.")
+
+    def wrapper(raw_f: Callable[..., Any]) -> _PartialFunction:
+        if isinstance(raw_f, _Function):
+            raw_f = raw_f.get_raw_f()
+            raise InvalidError(
+                f"Applying decorators for {raw_f} in the wrong order!\nUsage:\n\n"
+                "@app.function()\n@modal.batched()\ndef batched_function():\n    ..."
+            )
+        return _PartialFunction(
+            raw_f,
+            _PartialFunctionFlags.FUNCTION | _PartialFunctionFlags.BATCHED,
+            batch_max_size=max_batch_size,
+            batch_wait_ms=wait_ms,
+        )
+
+    return wrapper
+
+
 method = synchronize_api(_method)
 web_endpoint = synchronize_api(_web_endpoint)
 asgi_app = synchronize_api(_asgi_app)
@@ -512,3 +676,4 @@ web_server = synchronize_api(_web_server)
 build = synchronize_api(_build)
 enter = synchronize_api(_enter)
 exit = synchronize_api(_exit)
+batched = synchronize_api(_batched)