modAL 1.2.5__tar.gz → 1.2.5.dev1__tar.gz

{modal-1.2.5 → modal-1.2.5.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.2.5
+Version: 1.2.5.dev1
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.2.5 → modal-1.2.5.dev1}/modal/_container_entrypoint.py

@@ -15,19 +15,27 @@ if telemetry_socket:
     instrument_imports(telemetry_socket)
 
 import asyncio
+import inspect
 import queue
 import signal
+import sys
 import threading
 import time
-import types
-from typing import TYPE_CHECKING, Any, Optional, cast
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from google.protobuf.message import Message
 
 from modal._clustered_functions import initialize_clustered_function
-from modal._runtime.user_code_event_loop import UserCodeEventLoop
+from modal._partial_function import (
+    _find_callables_for_obj,
+    _PartialFunctionFlags,
+)
 from modal._serialization import deserialize, deserialize_params
 from modal._utils.async_utils import TaskContext, aclosing, synchronizer
+from modal._utils.function_utils import (
+    callable_has_non_self_params,
+)
 from modal.app import App, _App
 from modal.client import Client, _Client
 from modal.config import logger
@@ -45,7 +53,6 @@ from ._runtime.container_io_manager import (
 if TYPE_CHECKING:
     import modal._object
     import modal._runtime.container_io_manager
-    import modal._runtime.user_code_imports
 
 
 class DaemonizedThreadPool:
@@ -94,6 +101,81 @@ class DaemonizedThreadPool:
         self.inputs.put((func, args))
 
 
+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)
+
+
 def call_function(
     user_code_event_loop: UserCodeEventLoop,
     container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
@@ -249,25 +331,23 @@ def call_function(
                     signal.signal(signal.SIGUSR1, usr1_handler)  # reset signal handler
 
 
-def get_serialized_user_class_and_function(
-    function_def: api_pb2.Function, client: _Client
-) -> tuple[Optional[type], Optional[types.FunctionType]]:
-    if function_def.definition_type == api_pb2.Function.DEFINITION_TYPE_SERIALIZED:
-        assert function_def.function_serialized or function_def.class_serialized
-
-        if function_def.function_serialized:
-            ser_fun = deserialize(function_def.function_serialized, client)
-        else:
-            ser_fun = None
-
-        if function_def.class_serialized:
-            ser_usr_cls = deserialize(function_def.class_serialized, client)
-        else:
-            ser_usr_cls = None
-    else:
-        ser_usr_cls, ser_fun = None, None
-
-    return ser_usr_cls, ser_fun
+def call_lifecycle_functions(
+    event_loop: UserCodeEventLoop,
+    container_io_manager,  #: ContainerIOManager,  TODO: this type is generated at runtime
+    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.
+            # We can remove this once the deprecation in the actual @exit decorator is enforced.
+            args = (None, None, None) if callable_has_non_self_params(func) else ()
+            # in case func is non-async, it's executed here and sigint will by default
+            # interrupt it using a KeyboardInterrupt exception
+            res = func(*args)
+            if inspect.iscoroutine(res):
+                # if however func is async, we have to jump through some hoops
+                event_loop.run(res)
 
 
 def main(container_args: api_pb2.ContainerArguments, client: Client):
@@ -277,20 +357,34 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
     active_app: _App
     service: Service
     function_def = container_args.function_def
+    is_auto_snapshot: bool = function_def.is_auto_snapshot
     # The worker sets this flag to "1" for snapshot and restore tasks. Otherwise, this flag is unset,
     # in which case snapshots should be disabled.
     is_snapshotting_function = (
         function_def.is_checkpointing_function and os.environ.get("MODAL_ENABLE_SNAP_RESTORE") == "1"
     )
 
-    _client: _Client = cast(_Client, synchronizer._translate_in(client))  # TODO(erikbern): ugly
+    _client: _Client = synchronizer._translate_in(client)  # TODO(erikbern): ugly
 
     # Call ContainerHello - currently a noop but might be used later for things
     container_io_manager.hello()
 
     with container_io_manager.heartbeats(is_snapshotting_function), UserCodeEventLoop() as event_loop:
         # If this is a serialized function, fetch the definition from the server
-        ser_usr_cls, ser_fun = get_serialized_user_class_and_function(function_def, _client)
+        if function_def.definition_type == api_pb2.Function.DEFINITION_TYPE_SERIALIZED:
+            assert function_def.function_serialized or function_def.class_serialized
+
+            if function_def.function_serialized:
+                ser_fun = deserialize(function_def.function_serialized, _client)
+            else:
+                ser_fun = None
+
+            if function_def.class_serialized:
+                ser_usr_cls = deserialize(function_def.class_serialized, _client)
+            else:
+                ser_usr_cls = None
+        else:
+            ser_usr_cls, ser_fun = None, None
 
         # Initialize the function, importing user code.
         with container_io_manager.handle_user_exception():
@@ -342,7 +436,7 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
 
         # Initialize objects on the app.
         # This is basically only functions and classes - anything else is deprecated and will be unsupported soon
-        app: App = cast(App, synchronizer._translate_out(active_app))
+        app: App = synchronizer._translate_out(active_app)
         app._init_container(client, container_app)
 
         # Hydrate all function dependencies.
@@ -362,7 +456,7 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
                     "that evaluates differently in the local and remote environments."
                 )
             for object_id, obj in zip(dep_object_ids, service.service_deps):
-                metadata: Optional[Message] = container_app.object_handle_metadata[object_id]
+                metadata: Message = container_app.object_handle_metadata[object_id]
                 obj._hydrate(object_id, _client, metadata)
 
         # Initialize clustered functions.
@@ -373,8 +467,91 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
                 function_def._experimental_group_size,
             )
 
-        with service.execution_context(event_loop, container_io_manager) as finalized_functions:
-            call_function(event_loop, container_io_manager, finalized_functions, batch_max_size, batch_wait_ms)
+        # Identify all "enter" methods that need to run before we snapshot.
+        if service.user_cls_instance is not None and not is_auto_snapshot:
+            pre_snapshot_methods = _find_callables_for_obj(
+                service.user_cls_instance, _PartialFunctionFlags.ENTER_PRE_SNAPSHOT
+            )
+            call_lifecycle_functions(event_loop, container_io_manager, list(pre_snapshot_methods.values()))
+
+        # 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.
+        if is_snapshotting_function:
+            container_io_manager.memory_snapshot()
+
+        # 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
+
+            frame = inspect.currentframe().f_back
+
+            pdb.Pdb().set_trace(frame)
+
+        sys.breakpointhook = breakpoint_wrapper
+
+        # Identify the "enter" methods to run after resuming from a snapshot.
+        if service.user_cls_instance is not None and not is_auto_snapshot:
+            post_snapshot_methods = _find_callables_for_obj(
+                service.user_cls_instance, _PartialFunctionFlags.ENTER_POST_SNAPSHOT
+            )
+            call_lifecycle_functions(event_loop, container_io_manager, list(post_snapshot_methods.values()))
+
+        with container_io_manager.handle_user_exception():
+            finalized_functions = service.get_finalized_functions(function_def, container_io_manager)
+        # Execute the function.
+        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())
+            call_function(
+                event_loop,
+                container_io_manager,
+                finalized_functions,
+                batch_max_size,
+                batch_wait_ms,
+            )
+        finally:
+            # Run exit handlers. From this point onward, ignore all SIGINT signals that come from
+            # graceful shutdowns originating on the worker, as well as stray SIGUSR1 signals that
+            # may have been sent to cancel inputs.
+            int_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+            usr1_handler = signal.signal(signal.SIGUSR1, signal.SIG_IGN)
+
+            try:
+                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 lifespan_background_task in lifespan_background_tasks:
+                        lifespan_background_task.cancel()  # prevent dangling tasks
+
+                    # Identify "exit" methods and run them.
+                    # want to make sure this is called even if the lifespan manager fails
+                    if service.user_cls_instance is not None and not is_auto_snapshot:
+                        exit_methods = _find_callables_for_obj(service.user_cls_instance, _PartialFunctionFlags.EXIT)
+                        call_lifecycle_functions(event_loop, container_io_manager, list(exit_methods.values()))
+
+                # Finally, commit on exit to catch uncommitted volume changes and surface background
+                # commit errors.
+                container_io_manager.volume_commit(
+                    [v.volume_id for v in function_def.volume_mounts if v.allow_background_commits]
+                )
+            finally:
+                # Restore the original signal handler, needed for container_test hygiene since the
+                # test runs `main()` multiple times in the same process.
+                signal.signal(signal.SIGINT, int_handler)
+                signal.signal(signal.SIGUSR1, usr1_handler)
 
 
 if __name__ == "__main__":

{modal-1.2.5 → modal-1.2.5.dev1}/modal/_functions.py

@@ -701,7 +701,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         _experimental_proxy_ip: Optional[str] = None,
         _experimental_custom_scaling_factor: Optional[float] = None,
         restrict_output: bool = False,
-        http_config: Optional[api_pb2.HTTPConfig] = None,
     ) -> "_Function":
         """mdmd:hidden
 
@@ -988,7 +987,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 function_definition = api_pb2.Function(
                     module_name=info.module_name or "",
                     function_name=info.function_name,
-                    implementation_name=info.implementation_name,
                     mount_ids=loaded_mount_ids,
                     secret_ids=[secret.object_id for secret in secrets],
                     image_id=(image.object_id if image else ""),
@@ -1048,14 +1046,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     function_schema=function_schema,
                     supported_input_formats=supported_input_formats,
                     supported_output_formats=supported_output_formats,
-                    http_config=http_config,
                 )
 
                 if isinstance(gpu, list):
                     function_data = api_pb2.FunctionData(
                         module_name=function_definition.module_name,
                         function_name=function_definition.function_name,
-                        implementation_name=function_definition.implementation_name,
                         function_type=function_definition.function_type,
                         warm_pool_size=function_definition.warm_pool_size,
                         concurrency_limit=function_definition.concurrency_limit,
@@ -1087,7 +1083,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         untrusted=function_definition.untrusted,
                         supported_input_formats=supported_input_formats,
                         supported_output_formats=supported_output_formats,
-                        http_config=http_config,
                     )
 
                     ranked_functions = []

{modal-1.2.5 → modal-1.2.5.dev1}/modal/_partial_function.py

@@ -46,7 +46,6 @@ class _PartialFunctionFlags(enum.IntFlag):
     BATCHED = 64
     CONCURRENT = 128
     CLUSTERED = 256  # Experimental: Clustered functions
-    HTTP_WEB_INTERFACE = 512  # Experimental: HTTP server
 
     @staticmethod
     def all() -> int:
@@ -77,7 +76,6 @@ class _PartialFunctionParams:
     target_concurrent_inputs: Optional[int] = None
     build_timeout: Optional[int] = None
     rdma: Optional[bool] = None
-    http_config: Optional[api_pb2.HTTPConfig] = None
 
     def update(self, other: "_PartialFunctionParams") -> None:
         """Update self with params set in other."""