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

modal/_container_entrypoint.py

@@ -1,63 +1,103 @@
 # Copyright Modal Labs 2022
+# ruff: noqa: E402
+import os
+
+from modal._runtime.user_code_imports import Service, import_class_service, import_single_function_service
+
+telemetry_socket = os.environ.get("MODAL_TELEMETRY_SOCKET")
+if telemetry_socket:
+    from ._runtime.telemetry import instrument_imports
+
+    instrument_imports(telemetry_socket)
+
 import asyncio
-import base64
-import importlib
+import concurrent.futures
 import inspect
+import queue
 import signal
 import sys
 import threading
 import time
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Type
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from google.protobuf.message import Message
-from synchronicity import Interface
 
+from modal._clustered_functions import initialize_clustered_function
+from modal._proxy_tunnel import proxy_tunnel
+from modal._serialization import deserialize, deserialize_proto_params
+from modal._utils.async_utils import TaskContext, 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
+from modal.exception import ExecutionError, InputCancellation, InvalidError
+from modal.partial_function import (
+    _find_callables_for_obj,
+    _PartialFunctionFlags,
+)
+from modal.running_app import RunningApp
 from modal_proto import api_pb2
 
-from ._asgi import (
-    asgi_app_wrapper,
-    get_ip_address,
-    wait_for_web_server,
-    web_server_proxy,
-    webhook_asgi_app,
-    wsgi_app_wrapper,
+from ._runtime.container_io_manager import (
+    ContainerIOManager,
+    IOContext,
+    UserException,
+    _ContainerIOManager,
 )
-from ._container_io_manager import ContainerIOManager, UserException, _ContainerIOManager
-from ._proxy_tunnel import proxy_tunnel
-from ._serialization import deserialize
-from ._utils.async_utils import TaskContext, synchronizer
-from ._utils.function_utils import (
-    LocalFunctionError,
-    is_async as get_is_async,
-    is_global_function,
-    method_has_params,
-)
-from .app import App, _App
-from .client import Client, _Client
-from .cls import Cls
-from .config import logger
-from .exception import ExecutionError, InputCancellation, InvalidError
-from .execution_context import _set_current_context_ids, interact
-from .functions import Function, _Function
-from .partial_function import _find_callables_for_obj, _PartialFunctionFlags
-from .running_app import RunningApp
+from ._runtime.execution_context import _set_current_context_ids
 
 if TYPE_CHECKING:
-    from types import ModuleType
+    import modal._runtime.container_io_manager
+    import modal.object
 
 
-@dataclass
-class ImportedFunction:
-    obj: Any
-    fun: Callable
-    app: Optional[_App]
-    is_async: bool
-    is_generator: bool
-    data_format: int  # api_pb2.DataFormat
-    input_concurrency: int
-    is_auto_snapshot: bool
-    function: _Function
+class DaemonizedThreadPool:
+    # Used instead of ThreadPoolExecutor, since the latter won't allow
+    # the interpreter to shut down before the currently running tasks
+    # have finished
+    def __init__(self, max_threads: int):
+        self.max_threads = max_threads
+
+    def __enter__(self):
+        self.spawned_workers = 0
+        self.inputs: queue.Queue[Any] = queue.Queue()
+        self.finished = threading.Event()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.finished.set()
+
+        if exc_type is None:
+            self.inputs.join()
+        else:
+            # special case - allows us to exit the
+            if self.inputs.unfinished_tasks:
+                logger.info(
+                    f"Exiting DaemonizedThreadPool with {self.inputs.unfinished_tasks} active "
+                    f"inputs due to exception: {repr(exc_type)}"
+                )
+
+    def submit(self, func, *args):
+        def worker_thread():
+            while not self.finished.is_set():
+                try:
+                    _func, _args = self.inputs.get(timeout=1)
+                except queue.Empty:
+                    continue
+                try:
+                    _func(*_args)
+                except BaseException:
+                    logger.exception(f"Exception raised by {_func} in DaemonizedThreadPool worker!")
+                self.inputs.task_done()
+
+        if self.spawned_workers < self.max_threads:
+            threading.Thread(target=worker_thread, daemon=True).start()
+            self.spawned_workers += 1
+
+        self.inputs.put((func, args))
 
 
 class UserCodeEventLoop:
@@ -76,14 +116,25 @@ class UserCodeEventLoop:
 
     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
@@ -99,7 +150,9 @@ class UserCodeEventLoop:
                 # first sigint is graceful
                 task.cancel()
                 return
-            raise KeyboardInterrupt()  # this should normally not happen, but the second sigint would "hard kill" the event loop!
+
+            # 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:
@@ -122,111 +175,21 @@ class UserCodeEventLoop:
                 self.loop.remove_signal_handler(signal.SIGINT)
 
 
-def call_function_sync(
-    container_io_manager,  #: ContainerIOManager,  TODO: this type is generated at runtime
-    imp_fun: ImportedFunction,
+def call_function(
+    user_code_event_loop: UserCodeEventLoop,
+    container_io_manager: "modal._runtime.container_io_manager.ContainerIOManager",
+    finalized_functions: dict[str, "modal._runtime.user_code_imports.FinalizedFunction"],
+    batch_max_size: int,
+    batch_wait_ms: int,
 ):
-    def run_input(input_id: str, function_call_id: str, args: Any, kwargs: Any) -> None:
+    async def run_input_async(io_context: IOContext) -> None:
         started_at = time.time()
-        reset_context = _set_current_context_ids(input_id, function_call_id)
-        with container_io_manager.handle_input_exception(input_id, started_at):
-            logger.debug(f"Starting input {input_id} (sync)")
-            res = imp_fun.fun(*args, **kwargs)
-            logger.debug(f"Finished input {input_id} (sync)")
-
+        input_ids, function_call_ids = io_context.input_ids, io_context.function_call_ids
+        reset_context = _set_current_context_ids(input_ids, function_call_ids)
+        async with container_io_manager.handle_input_exception.aio(io_context, started_at):
+            res = io_context.call_finalized_function()
             # TODO(erikbern): any exception below shouldn't be considered a user exception
-            if imp_fun.is_generator:
-                if not inspect.isgenerator(res):
-                    raise InvalidError(f"Generator function returned value of type {type(res)}")
-
-                # Send up to this many outputs at a time.
-                generator_queue: asyncio.Queue[Any] = container_io_manager._queue_create(1024)
-                generator_output_task = container_io_manager.generator_output_task(
-                    function_call_id,
-                    imp_fun.data_format,
-                    generator_queue,
-                    _future=True,  # Synchronicity magic to return a future.
-                )
-
-                item_count = 0
-                for value in res:
-                    container_io_manager._queue_put(generator_queue, value)
-                    item_count += 1
-
-                container_io_manager._queue_put(generator_queue, _ContainerIOManager._GENERATOR_STOP_SENTINEL)
-                generator_output_task.result()  # Wait to finish sending generator outputs.
-                message = api_pb2.GeneratorDone(items_total=item_count)
-                container_io_manager.push_output(input_id, started_at, message, api_pb2.DATA_FORMAT_GENERATOR_DONE)
-            else:
-                if inspect.iscoroutine(res) or inspect.isgenerator(res) or inspect.isasyncgen(res):
-                    raise InvalidError(
-                        f"Sync (non-generator) function return value of type {type(res)}."
-                        " You might need to use @app.function(..., is_generator=True)."
-                    )
-                container_io_manager.push_output(input_id, started_at, res, imp_fun.data_format)
-        reset_context()
-
-    if imp_fun.input_concurrency > 1:
-        # We can't use `concurrent.futures.ThreadPoolExecutor` here because in Python 3.11+, this
-        # class has no workaround that allows us to exit the Python interpreter process without
-        # waiting for the worker threads to finish. We need this behavior on SIGINT.
-
-        import queue
-        import threading
-
-        spawned_workers = 0
-        inputs: queue.Queue[Any] = queue.Queue()
-        finished = threading.Event()
-
-        def worker_thread():
-            while not finished.is_set():
-                try:
-                    args = inputs.get(timeout=1)
-                except queue.Empty:
-                    continue
-                try:
-                    run_input(*args)
-                except BaseException:
-                    # This should basically never happen, since only KeyboardInterrupt is the only error that can
-                    # bubble out of from handle_input_exception and those wouldn't be raised outside the main thread
-                    pass
-                inputs.task_done()
-
-        for input_id, function_call_id, args, kwargs in container_io_manager.run_inputs_outputs(
-            imp_fun.input_concurrency
-        ):
-            if spawned_workers < imp_fun.input_concurrency:
-                threading.Thread(target=worker_thread, daemon=True).start()
-                spawned_workers += 1
-            inputs.put((input_id, function_call_id, args, kwargs))
-
-        finished.set()
-        inputs.join()
-
-    else:
-        for input_id, function_call_id, args, kwargs in container_io_manager.run_inputs_outputs(
-            imp_fun.input_concurrency
-        ):
-            try:
-                run_input(input_id, function_call_id, args, kwargs)
-            except:
-                raise
-
-
-async def call_function_async(
-    container_io_manager,  #: ContainerIOManager,  TODO: this type is generated at runtime
-    imp_fun: ImportedFunction,
-):
-    async def run_input(input_id: str, function_call_id: str, args: Any, kwargs: Any) -> None:
-        started_at = time.time()
-        reset_context = _set_current_context_ids(input_id, function_call_id)
-        async with container_io_manager.handle_input_exception.aio(input_id, started_at):
-            logger.debug(f"Starting input {input_id} (async)")
-            res = imp_fun.fun(*args, **kwargs)
-            logger.debug(f"Finished input {input_id} (async)")
-
-            # TODO(erikbern): any exception below shouldn't be considered a user exception
-            if imp_fun.is_generator:
+            if io_context.finalized_function.is_generator:
                 if not inspect.isasyncgen(res):
                     raise InvalidError(f"Async generator function returned value of type {type(res)}")
 
@@ -234,8 +197,8 @@ async def call_function_async(
                 generator_queue: asyncio.Queue[Any] = await container_io_manager._queue_create.aio(1024)
                 generator_output_task = asyncio.create_task(
                     container_io_manager.generator_output_task.aio(
-                        function_call_id,
-                        imp_fun.data_format,
+                        function_call_ids[0],
+                        io_context.finalized_function.data_format,
                         generator_queue,
                     )
                 )
@@ -248,8 +211,11 @@ async def call_function_async(
                 await container_io_manager._queue_put.aio(generator_queue, _ContainerIOManager._GENERATOR_STOP_SENTINEL)
                 await generator_output_task  # Wait to finish sending generator outputs.
                 message = api_pb2.GeneratorDone(items_total=item_count)
-                await container_io_manager.push_output.aio(
-                    input_id, started_at, message, api_pb2.DATA_FORMAT_GENERATOR_DONE
+                await container_io_manager.push_outputs.aio(
+                    io_context,
+                    started_at,
+                    message,
+                    api_pb2.DATA_FORMAT_GENERATOR_DONE,
                 )
             else:
                 if not inspect.iscoroutine(res) or inspect.isgenerator(res) or inspect.isasyncgen(res):
@@ -258,310 +224,332 @@ async def call_function_async(
                         " You might need to use @app.function(..., is_generator=True)."
                     )
                 value = await res
-                await container_io_manager.push_output.aio(input_id, started_at, value, imp_fun.data_format)
+                await container_io_manager.push_outputs.aio(
+                    io_context,
+                    started_at,
+                    value,
+                    io_context.finalized_function.data_format,
+                )
         reset_context()
 
-    if imp_fun.input_concurrency > 1:
-        # all run_input coroutines will have completed by the time we leave the execution context
-        # but the wrapping *tasks* may not yet have been resolved, so we add a 0.01s
-        # for them to resolve gracefully:
-        async with TaskContext(0.01) as task_context:
-            async for input_id, function_call_id, args, kwargs in container_io_manager.run_inputs_outputs.aio(
-                imp_fun.input_concurrency
-            ):
-                # Note that run_inputs_outputs will not return until the concurrency semaphore has
-                # released all its slots so that they can be acquired by the run_inputs_outputs finalizer
-                # This prevents leaving the task_context before outputs have been created
-                # TODO: refactor to make this a bit more easy to follow?
-                task_context.create_task(run_input(input_id, function_call_id, args, kwargs))
-    else:
-        async for input_id, function_call_id, args, kwargs in container_io_manager.run_inputs_outputs.aio(
-            imp_fun.input_concurrency
-        ):
-            await run_input(input_id, function_call_id, args, kwargs)
-
-
-def import_function(
-    function_def: api_pb2.Function,
-    ser_cls,
-    ser_fun,
-    ser_params: Optional[bytes],
-    container_io_manager,
-    client: Client,
-) -> ImportedFunction:
-    """Imports a function dynamically, and locates the app.
-
-    This is somewhat complex because we're dealing with 3 quite different type of functions:
-    1. Functions defined in global scope and decorated in global scope (Function objects)
-    2. Functions defined in global scope but decorated elsewhere (these will be raw callables)
-    3. Serialized functions
-
-    In addition, we also need to handle
-    * Normal functions
-    * Methods on classes (in which case we need to instantiate the object)
-
-    This helper also handles web endpoints, ASGI/WSGI servers, and HTTP servers.
-
-    In order to locate the app, we try two things:
-    * If the function is a Function, we can get the app directly from it
-    * Otherwise, use the app name and look it up from a global list of apps: this
-      typically only happens in case 2 above, or in sometimes for case 3
-
-    Note that `import_function` is *not* synchronized, becase we need it to run on the main
-    thread. This is so that any user code running in global scope (which executes as a part of
-    the import) runs on the right thread.
-    """
-    module: Optional[ModuleType] = None
-    cls: Optional[Type] = None
-    fun: Callable
-    function: Optional[_Function] = None
-    active_app: Optional[_App] = None
-    pty_info: api_pb2.PTYInfo = function_def.pty_info
-
-    if ser_fun is not None:
-        # This is a serialized function we already fetched from the server
-        cls, fun = ser_cls, ser_fun
-    else:
-        # Load the module dynamically
-        module = importlib.import_module(function_def.module_name)
-        qual_name: str = function_def.function_name
-
-        if not is_global_function(qual_name):
-            raise LocalFunctionError("Attempted to load a function defined in a function scope")
-
-        parts = qual_name.split(".")
-        if len(parts) == 1:
-            # This is a function
-            cls = None
-            f = getattr(module, qual_name)
-            if isinstance(f, Function):
-                function = synchronizer._translate_in(f)
-                fun = function.get_raw_f()
-                active_app = function._app
-            else:
-                fun = f
-        elif len(parts) == 2:
-            # This is a method on a class
-            cls_name, fun_name = parts
-            cls = getattr(module, cls_name)
-            if isinstance(cls, Cls):
-                # The cls decorator is in global scope
-                _cls = synchronizer._translate_in(cls)
-                fun = _cls._callables[fun_name]
-                function = _cls._functions.get(fun_name)
-                active_app = _cls._app
-            else:
-                # This is a raw class
-                fun = getattr(cls, fun_name)
-        else:
-            raise InvalidError(f"Invalid function qualname {qual_name}")
-
-    # If the cls/function decorator was applied in local scope, but the app is global, we can look it up
-    if active_app is None:
-        # This branch is reached in the special case that the imported function is 1) not serialized, and 2) isn't a FunctionHandle - i.e, not decorated at definition time
-        # Look at all instantiated apps - if there is only one with the indicated name, use that one
-        app_name: Optional[str] = function_def.app_name or None  # coalesce protobuf field to None
-        matching_apps = _App._all_apps.get(app_name, [])
-        if len(matching_apps) > 1:
-            if app_name is not None:
-                warning_sub_message = f"app with the same name ('{app_name}')"
-            else:
-                warning_sub_message = "unnamed app"
-            logger.warning(
-                f"You have more than one {warning_sub_message}. It's recommended to name all your Apps uniquely when using multiple apps"
-            )
-        elif len(matching_apps) == 1:
-            (active_app,) = matching_apps
-        # there could also technically be zero found apps, but that should probably never be an issue since that would mean user won't use is_inside or other function handles anyway
+    def run_input_sync(io_context: IOContext) -> None:
+        started_at = time.time()
+        input_ids, function_call_ids = io_context.input_ids, io_context.function_call_ids
+        reset_context = _set_current_context_ids(input_ids, function_call_ids)
+        with container_io_manager.handle_input_exception(io_context, started_at):
+            res = io_context.call_finalized_function()
 
-    # Check this property before we turn it into a method (overriden by webhooks)
-    is_async = get_is_async(fun)
+            # TODO(erikbern): any exception below shouldn't be considered a user exception
+            if io_context.finalized_function.is_generator:
+                if not inspect.isgenerator(res):
+                    raise InvalidError(f"Generator function returned value of type {type(res)}")
 
-    # Use the function definition for whether this is a generator (overriden by webhooks)
-    is_generator = function_def.function_type == api_pb2.Function.FUNCTION_TYPE_GENERATOR
+                # Send up to this many outputs at a time.
+                generator_queue: asyncio.Queue[Any] = container_io_manager._queue_create(1024)
+                generator_output_task: concurrent.futures.Future = container_io_manager.generator_output_task(  # type: ignore
+                    function_call_ids[0],
+                    io_context.finalized_function.data_format,
+                    generator_queue,
+                    _future=True,  # type: ignore  # Synchronicity magic to return a future.
+                )
 
-    # What data format is used for function inputs and outputs
-    data_format = api_pb2.DATA_FORMAT_PICKLE
+                item_count = 0
+                for value in res:
+                    container_io_manager._queue_put(generator_queue, value)
+                    item_count += 1
 
-    # Container can fetch multiple inputs simultaneously
-    if pty_info.pty_type == api_pb2.PTYInfo.PTY_TYPE_SHELL:
-        # Concurrency doesn't apply for `modal shell`.
-        input_concurrency = 1
-    else:
-        input_concurrency = function_def.allow_concurrent_inputs or 1
+                container_io_manager._queue_put(generator_queue, _ContainerIOManager._GENERATOR_STOP_SENTINEL)
+                generator_output_task.result()  # Wait to finish sending generator outputs.
+                message = api_pb2.GeneratorDone(items_total=item_count)
+                container_io_manager.push_outputs(io_context, started_at, message, api_pb2.DATA_FORMAT_GENERATOR_DONE)
+            else:
+                if inspect.iscoroutine(res) or inspect.isgenerator(res) or inspect.isasyncgen(res):
+                    raise InvalidError(
+                        f"Sync (non-generator) function return value of type {type(res)}."
+                        " You might need to use @app.function(..., is_generator=True)."
+                    )
+                container_io_manager.push_outputs(
+                    io_context, started_at, res, io_context.finalized_function.data_format
+                )
+        reset_context()
 
-    # Instantiate the class if it's defined
-    if cls:
-        if ser_params:
-            _client: _Client = synchronizer._translate_in(client)
-            args, kwargs = deserialize(ser_params, _client)
-        else:
-            args, kwargs = (), {}
-        obj = cls(*args, **kwargs)
-        if isinstance(cls, Cls):
-            obj = obj.get_obj()
-        # Bind the function to the instance (using the descriptor protocol!)
-        fun = fun.__get__(obj)
+    if container_io_manager.target_concurrency > 1:
+        with DaemonizedThreadPool(max_threads=container_io_manager.max_concurrency) as thread_pool:
+
+            def make_async_cancel_callback(task):
+                def f():
+                    user_code_event_loop.loop.call_soon_threadsafe(task.cancel)
+
+                return f
+
+            did_sigint = False
+
+            def cancel_callback_sync():
+                nonlocal did_sigint
+                # We only want one sigint even if multiple inputs are cancelled
+                # A second sigint would forcibly shut down the event loop and spew
+                # out a bunch of tracebacks, which we only want to happen in case
+                # the worker kills this process after a failed self-termination
+                if not did_sigint:
+                    did_sigint = True
+                    logger.warning(
+                        "User cancelling input of non-async functions with allow_concurrent_inputs > 1.\n"
+                        "This shuts down the container, causing concurrently running inputs to be "
+                        "rescheduled in other containers."
+                    )
+                    os.kill(os.getpid(), signal.SIGINT)
+
+            async def run_concurrent_inputs():
+                # all run_input coroutines will have completed by the time we leave the execution context
+                # but the wrapping *tasks* may not yet have been resolved, so we add a 0.01s
+                # for them to resolve gracefully:
+                async with TaskContext(0.01) as task_context:
+                    async for io_context in container_io_manager.run_inputs_outputs.aio(
+                        finalized_functions, batch_max_size, batch_wait_ms
+                    ):
+                        # Note that run_inputs_outputs will not return until all the input slots are released
+                        # so that they can be acquired by the run_inputs_outputs finalizer
+                        # This prevents leaving the task_context before outputs have been created
+                        # TODO: refactor to make this a bit more easy to follow?
+                        if io_context.finalized_function.is_async:
+                            input_task = task_context.create_task(run_input_async(io_context))
+                            io_context.set_cancel_callback(make_async_cancel_callback(input_task))
+                        else:
+                            # run sync input in thread
+                            thread_pool.submit(run_input_sync, io_context)
+                            io_context.set_cancel_callback(cancel_callback_sync)
+
+            user_code_event_loop.run(run_concurrent_inputs())
     else:
-        obj = None
-
-    if function_def.webhook_config.type:
-        is_async = True
-        is_generator = True
-        data_format = api_pb2.DATA_FORMAT_ASGI
-
-        if function_def.webhook_config.type == api_pb2.WEBHOOK_TYPE_ASGI_APP:
-            # Function returns an asgi_app, which we can use as a callable.
-            fun = asgi_app_wrapper(fun(), container_io_manager)
-
-        elif function_def.webhook_config.type == api_pb2.WEBHOOK_TYPE_WSGI_APP:
-            # Function returns an wsgi_app, which we can use as a callable.
-            fun = wsgi_app_wrapper(fun(), container_io_manager)
-
-        elif function_def.webhook_config.type == api_pb2.WEBHOOK_TYPE_FUNCTION:
-            # Function is a webhook without an ASGI app. Create one for it.
-            fun = asgi_app_wrapper(
-                webhook_asgi_app(fun, function_def.webhook_config.method),
-                container_io_manager,
-            )
-
-        elif function_def.webhook_config.type == api_pb2.WEBHOOK_TYPE_WEB_SERVER:
-            # Function spawns an HTTP web server listening at a port.
-            fun()
-
-            # We intentionally try to connect to the external interface instead of the loopback
-            # interface here so users are forced to expose the server. This allows us to potentially
-            # change the implementation to use an external bridge in the future.
-            host = get_ip_address(b"eth0")
-            port = function_def.webhook_config.web_server_port
-            startup_timeout = function_def.webhook_config.web_server_startup_timeout
-            wait_for_web_server(host, port, timeout=startup_timeout)
-            fun = asgi_app_wrapper(web_server_proxy(host, port), container_io_manager)
+        for io_context in container_io_manager.run_inputs_outputs(finalized_functions, batch_max_size, batch_wait_ms):
+            if io_context.finalized_function.is_async:
+                user_code_event_loop.run(run_input_async(io_context))
+            else:
+                # Set up a custom signal handler for `SIGUSR1`, which gets translated to an InputCancellation
+                # during function execution. This is sent to cancel inputs from the user
+                def _cancel_input_signal_handler(signum, stackframe):
+                    raise InputCancellation("Input was cancelled by user")
 
+                usr1_handler = signal.signal(signal.SIGUSR1, _cancel_input_signal_handler)
+                # run this sync code in the main thread, blocking the "userland" event loop
+                # this lets us cancel it using a signal handler that raises an exception
+                try:
+                    run_input_sync(io_context)
+                finally:
+                    signal.signal(signal.SIGUSR1, usr1_handler)  # reset signal handler
+
+
+def get_active_app_fallback(function_def: api_pb2.Function) -> _App:
+    # This branch is reached in the special case that the imported function/class is:
+    # 1) not serialized, and
+    # 2) isn't a FunctionHandle - i.e, not decorated at definition time
+    # Look at all instantiated apps - if there is only one with the indicated name, use that one
+    app_name: Optional[str] = function_def.app_name or None  # coalesce protobuf field to None
+    matching_apps = _App._all_apps.get(app_name, [])
+    if len(matching_apps) == 1:
+        active_app: _App = matching_apps[0]
+        return active_app
+
+    if len(matching_apps) > 1:
+        if app_name is not None:
+            warning_sub_message = f"app with the same name ('{app_name}')"
         else:
-            raise InvalidError(f"Unrecognized web endpoint type {function_def.webhook_config.type}")
-
-    return ImportedFunction(
-        obj,
-        fun,
-        active_app,
-        is_async,
-        is_generator,
-        data_format,
-        input_concurrency,
-        function_def.is_auto_snapshot,
-        function,
-    )
+            warning_sub_message = "unnamed app"
+        logger.warning(
+            f"You have more than one {warning_sub_message}. "
+            "It's recommended to name all your Apps uniquely when using multiple apps"
+        )
+
+    # If we don't have an active app, create one on the fly
+    # The app object is used to carry the app layout etc
+    return _App()
 
 
 def call_lifecycle_functions(
     event_loop: UserCodeEventLoop,
     container_io_manager,  #: ContainerIOManager,  TODO: this type is generated at runtime
-    funcs: Sequence[Callable],
+    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 parameterized 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 method_has_params(func) else ()
-            res = func(
-                *args
-            )  # in case func is non-async, it's executed here and sigint will by default interrupt it using a KeyboardInterrupt exception
+            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 deserialize_params(serialized_params: bytes, function_def: api_pb2.Function, _client: "modal.client._Client"):
+    if function_def.class_parameter_info.format in (
+        api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_UNSPECIFIED,
+        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)
+    elif function_def.class_parameter_info.format == api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PROTO:
+        param_args = ()
+        param_kwargs = deserialize_proto_params(serialized_params, list(function_def.class_parameter_info.schema))
+    else:
+        raise ExecutionError(
+            f"Unknown class parameter serialization format: {function_def.class_parameter_info.format}"
+        )
+
+    return param_args, param_kwargs
+
+
 def main(container_args: api_pb2.ContainerArguments, client: Client):
     # This is a bit weird but we need both the blocking and async versions of ContainerIOManager.
     # At some point, we should fix that by having built-in support for running "user code"
     container_io_manager = ContainerIOManager(container_args, 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", "0") == "1"
+    )
+
+    _client: _Client = synchronizer._translate_in(client)  # TODO(erikbern): ugly
 
-    with container_io_manager.heartbeats(), UserCodeEventLoop() as event_loop:
+    # 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
-        if container_args.function_def.definition_type == api_pb2.Function.DEFINITION_TYPE_SERIALIZED:
+        if function_def.definition_type == api_pb2.Function.DEFINITION_TYPE_SERIALIZED:
             ser_cls, ser_fun = container_io_manager.get_serialized_function()
         else:
             ser_cls, ser_fun = None, None
 
         # Initialize the function, importing user code.
         with container_io_manager.handle_user_exception():
-            imp_fun = import_function(
-                container_args.function_def,
-                ser_cls,
-                ser_fun,
-                container_args.serialized_params,
-                container_io_manager,
-                client,
-            )
+            if container_args.serialized_params:
+                param_args, param_kwargs = deserialize_params(container_args.serialized_params, function_def, _client)
+            else:
+                param_args = ()
+                param_kwargs = {}
+
+            if function_def.is_class:
+                service = import_class_service(
+                    function_def,
+                    ser_cls,
+                    param_args,
+                    param_kwargs,
+                )
+            else:
+                service = import_single_function_service(
+                    function_def,
+                    ser_cls,
+                    ser_fun,
+                    param_args,
+                    param_kwargs,
+                )
+
+            # If the cls/function decorator was applied in local scope, but the app is global, we can look it up
+            if service.app is not None:
+                active_app = service.app
+            else:
+                # if the app can't be inferred by the imported function, use name-based fallback
+                active_app = get_active_app_fallback(function_def)
+
+            if function_def.pty_info.pty_type == api_pb2.PTYInfo.PTY_TYPE_SHELL:
+                # Concurrency and batching doesn't apply for `modal shell`.
+                batch_max_size = 0
+                batch_wait_ms = 0
+            else:
+                batch_max_size = function_def.batch_max_size or 0
+                batch_wait_ms = function_def.batch_linger_ms or 0
 
         # Get ids and metadata for objects (primarily functions and classes) on the app
-        container_app: RunningApp = container_io_manager.get_app_objects()
+        container_app: RunningApp = container_io_manager.get_app_objects(container_args.app_layout)
 
         # Initialize objects on the app.
         # This is basically only functions and classes - anything else is deprecated and will be unsupported soon
-        if imp_fun.app is not None:
-            app: App = synchronizer._translate_out(imp_fun.app, Interface.BLOCKING)
-            app._init_container(client, container_app)
+        app: App = synchronizer._translate_out(active_app)
+        app._init_container(client, container_app)
 
         # Hydrate all function dependencies.
         # TODO(erikbern): we an remove this once we
         # 1. Enable lazy hydration for all objects
         # 2. Fully deprecate .new() objects
-        if imp_fun.function:
-            _client: _Client = synchronizer._translate_in(client)  # TODO(erikbern): ugly
-            dep_object_ids: List[str] = [dep.object_id for dep in container_args.function_def.object_dependencies]
-            function_deps = imp_fun.function.deps(only_explicit_mounts=True)
-            if len(function_deps) != len(dep_object_ids):
+        if service.code_deps is not None:  # this is not set for serialized or non-global scope functions
+            dep_object_ids: list[str] = [dep.object_id for dep in function_def.object_dependencies]
+            if len(service.code_deps) != len(dep_object_ids):
                 raise ExecutionError(
-                    f"Function has {len(function_deps)} dependencies"
-                    f" but container got {len(dep_object_ids)} object ids."
+                    f"Function has {len(service.code_deps)} dependencies"
+                    f" but container got {len(dep_object_ids)} object ids.\n"
+                    f"Code deps: {service.code_deps}\n"
+                    f"Object ids: {dep_object_ids}"
                 )
-            for object_id, obj in zip(dep_object_ids, function_deps):
+            for object_id, obj in zip(dep_object_ids, service.code_deps):
                 metadata: Message = container_app.object_handle_metadata[object_id]
                 obj._hydrate(object_id, _client, metadata)
 
-        # Identify all "enter" methods that need to run before we checkpoint.
-        if imp_fun.obj is not None and not imp_fun.is_auto_snapshot:
-            pre_checkpoint_methods = _find_callables_for_obj(imp_fun.obj, _PartialFunctionFlags.ENTER_PRE_CHECKPOINT)
-            call_lifecycle_functions(event_loop, container_io_manager, list(pre_checkpoint_methods.values()))
+        # Initialize clustered functions.
+        if function_def._experimental_group_size > 0:
+            initialize_clustered_function(
+                client,
+                container_args.task_id,
+                function_def._experimental_group_size,
+            )
+
+        # 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 innitialization. Checkpointed containers run from this point onwards.
-        if container_args.function_def.is_checkpointing_function:
-            container_io_manager.checkpoint()
+        # 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.
-        if container_args.function_def.pty_info.pty_type != api_pb2.PTYInfo.PTY_TYPE_UNSPECIFIED:
+        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
 
-            def breakpoint_wrapper():
-                # note: it would be nice to not have breakpoint_wrapper() included in the backtrace
-                interact()
-                import pdb
+            frame = inspect.currentframe().f_back
 
-                pdb.set_trace()
+            pdb.Pdb().set_trace(frame)
 
-            sys.breakpointhook = breakpoint_wrapper
+        sys.breakpointhook = breakpoint_wrapper
 
-        # Identify the "enter" methods to run after resuming from a checkpoint.
-        if imp_fun.obj is not None and not imp_fun.is_auto_snapshot:
-            post_checkpoint_methods = _find_callables_for_obj(imp_fun.obj, _PartialFunctionFlags.ENTER_POST_CHECKPOINT)
-            call_lifecycle_functions(event_loop, container_io_manager, list(post_checkpoint_methods.values()))
+        # 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:
-            if imp_fun.is_async:
-                event_loop.run(call_function_async(container_io_manager, imp_fun))
-            else:
-                # Set up a signal handler for `SIGUSR1`, which gets translated to an InputCancellation
-                # during function execution. This is sent to cancel inputs from the user.
-                def _cancel_input_signal_handler(signum, stackframe):
-                    raise InputCancellation("Input was cancelled by user")
-
-                signal.signal(signal.SIGUSR1, _cancel_input_signal_handler)
-
-                call_function_sync(container_io_manager, imp_fun)
+            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
@@ -570,15 +558,27 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
             usr1_handler = signal.signal(signal.SIGUSR1, signal.SIG_IGN)
 
             try:
-                # Identify "exit" methods and run them.
-                if imp_fun.obj is not None and not imp_fun.is_auto_snapshot:
-                    exit_methods = _find_callables_for_obj(imp_fun.obj, _PartialFunctionFlags.EXIT)
-                    call_lifecycle_functions(event_loop, container_io_manager, list(exit_methods.values()))
+                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 container_args.function_def.volume_mounts if v.allow_background_commits]
+                    [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
@@ -591,7 +591,15 @@ if __name__ == "__main__":
     logger.debug("Container: starting")
 
     container_args = api_pb2.ContainerArguments()
-    container_args.ParseFromString(base64.b64decode(sys.argv[1]))
+
+    container_arguments_path: Optional[str] = os.environ.get("MODAL_CONTAINER_ARGUMENTS_PATH")
+    if container_arguments_path is None:
+        # TODO(erikbern): this fallback is for old workers and we can remove it very soon (days)
+        import base64
+
+        container_args.ParseFromString(base64.b64decode(sys.argv[1]))
+    else:
+        container_args.ParseFromString(open(container_arguments_path, "rb").read())
 
     # Note that we're creating the client in a synchronous context, but it will be running in a separate thread.
     # This is good because if the function is long running then we the client can still send heartbeats
@@ -611,7 +619,7 @@ if __name__ == "__main__":
     # from shutting down. The sleep(0) here is needed for finished ThreadPoolExecutor resources to
     # shut down without triggering this warning (e.g., `@wsgi_app()`).
     time.sleep(0)
-    lingering_threads: List[threading.Thread] = []
+    lingering_threads: list[threading.Thread] = []
     for thread in threading.enumerate():
         current_thread = threading.get_ident()
         if thread.ident is not None and thread.ident != current_thread and not thread.daemon and thread.is_alive():
@@ -619,7 +627,8 @@ if __name__ == "__main__":
     if lingering_threads:
         thread_names = ", ".join(t.name for t in lingering_threads)
         logger.warning(
-            f"Detected {len(lingering_threads)} background thread(s) [{thread_names}] still running after container exit. This will prevent runner shutdown for up to 30 seconds."
+            f"Detected {len(lingering_threads)} background thread(s) [{thread_names}] still running "
+            "after container exit. This will prevent runner shutdown for up to 30 seconds."
         )
 
     logger.debug("Container: done")