modal 1.4.4.dev14__tar.gz → 1.4.4.dev16__tar.gz

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/PKG-INFO

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

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/_container_entrypoint.py

@@ -14,7 +14,6 @@ if telemetry_socket:
     from ._runtime.telemetry import instrument_imports
 
     instrument_imports(telemetry_socket)
-
 import asyncio
 import queue
 import signal
@@ -41,10 +40,11 @@ from ._runtime.container_io_manager import (
     ContainerIOManager,
     IOContext,
 )
-from ._runtime.task_lifecycle_manager import UserException
+from ._runtime.task_lifecycle_manager import TaskLifecycleManager, UserException
 
 if TYPE_CHECKING:
     import modal._runtime.container_io_manager
+    import modal._runtime.task_lifecycle_manager
     import modal._runtime.user_code_imports
 
 
@@ -287,12 +287,146 @@ def get_serialized_user_class_and_function(
     return ser_usr_cls, ser_fun
 
 
-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)
+def get_server_service(
+    container_args: api_pb2.ContainerArguments, function_def: api_pb2.Function, _client: _Client, client: Client
+) -> Service:
+    ser_usr_cls, ser_fun = get_serialized_user_class_and_function(function_def, _client)
+    service_base_function_id = container_args.app_layout.function_ids[function_def.function_name]
+    service_function_hydration_data = [
+        o for o in container_args.app_layout.objects if o.object_id == service_base_function_id
+    ][0]
+    service = import_server_service(
+        function_def,
+        service_function_hydration_data,
+        client,
+        ser_usr_cls,
+    )
+    return service
+
+
+def get_class_service(
+    container_args: api_pb2.ContainerArguments, function_def: api_pb2.Function, _client: _Client, client: Client
+) -> Service:
+    ser_usr_cls, ser_fun = get_serialized_user_class_and_function(function_def, _client)
+    if container_args.serialized_params:
+        param_args, param_kwargs = deserialize_params(container_args.serialized_params, function_def, _client)
+    else:
+        param_args = ()
+        param_kwargs = {}
+    # this is a bit ugly - match the function and class based on function name to get metadata
+    # This metadata is required in order to hydrate the class in case it's not globally
+    # decorated (or serialized)
+    service_base_function_id = container_args.app_layout.function_ids[function_def.function_name]
+    service_function_hydration_data = [
+        o for o in container_args.app_layout.objects if o.object_id == service_base_function_id
+    ][0]
+    class_id = container_args.app_layout.class_ids[function_def.function_name.removesuffix(".*")]
+
+    service = import_class_service(
+        function_def,
+        service_function_hydration_data,
+        class_id,
+        client,
+        ser_usr_cls,
+        param_args,
+        param_kwargs,
+    )
+    return service
+
+
+def get_function_service(function_def: api_pb2.Function, _client: _Client) -> Service:
+    ser_usr_cls, ser_fun = get_serialized_user_class_and_function(function_def, _client)
+    assert ser_usr_cls is None
+    service = import_single_function_service(
+        function_def,
+        ser_fun,
+    )
+    return service
+
+
+def hydrate_function(
+    container_args: api_pb2.ContainerArguments,
+    task_lifecycle_manager: TaskLifecycleManager,
+    function_def: api_pb2.Function,
+    _client: _Client,
+    client: Client,
+) -> Service:
     active_app: _App
     service: Service
+    with task_lifecycle_manager.handle_task_lifecycle_exception():
+        with execution_context._import_context():
+            if function_def.is_server:
+                service = get_server_service(container_args, function_def, _client, client)
+            elif function_def.is_class:
+                service = get_class_service(container_args, function_def, _client, client)
+            else:
+                service = get_function_service(function_def, _client)
+
+        active_app = service.app
+
+    # Get ids and metadata for objects (primarily functions and classes) on the app
+    container_app: RunningApp = running_app_from_layout(container_args.app_id, 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
+    app: App = cast(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 service.service_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.service_deps) != len(dep_object_ids):
+            raise ExecutionError(
+                f"Function has {len(service.service_deps)} dependencies"
+                f" but container got {len(dep_object_ids)} object ids.\n"
+                f"Code deps: {service.service_deps}\n"
+                f"Object ids: {dep_object_ids}\n"
+                "\n"
+                "This can happen if you are defining Modal objects under a conditional statement "
+                "that evaluates differently in the local and remote environments."
+            )
+        for object_id, obj in zip(dep_object_ids, service.service_deps):
+            metadata: Message | None = container_app.object_handle_metadata[object_id]
+            obj._hydrate(object_id, _client, metadata)
+
+    return service
+
+
+def run_server(container_args: api_pb2.ContainerArguments, client: Client):
+    function_def = container_args.function_def
+    _client: _Client = cast(_Client, synchronizer._translate_in(client))  # TODO(erikbern): ugly
+    task_lifecycle_manager = TaskLifecycleManager(
+        container_args.task_id,
+        container_args.function_id,
+        container_args.function_def,
+        container_args.checkpoint_id or None,
+        client,
+    )
+
+    service: Service = hydrate_function(container_args, task_lifecycle_manager, function_def, _client, client)
+    # Initialize clustered functions.
+    if function_def._experimental_group_size > 0:
+        initialize_clustered_function(
+            client,
+            container_args.task_id,
+            function_def._experimental_group_size,
+        )
+
+    with UserCodeEventLoop() as event_loop:
+        with service.lifecycle_context(
+            event_loop,
+            task_lifecycle_manager=task_lifecycle_manager,
+        ):
+            call_server(event_loop)
+
+
+def run_function(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: Any = ContainerIOManager(container_args, client)
     function_def = container_args.function_def
     # The worker sets this flag to "1" for snapshot and restore tasks. Otherwise, this flag is unset,
     # in which case snapshots should be disabled.
@@ -306,92 +440,17 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
     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)
-
         # Initialize the function, importing user code.
-        with container_io_manager.get_task_lifecycle_manager().handle_task_lifecycle_exception():
-            if container_args.serialized_params:
-                param_args, param_kwargs = deserialize_params(container_args.serialized_params, function_def, _client)
-            else:
-                param_args = ()
-                param_kwargs = {}
-
-            with execution_context._import_context():
-                if function_def.is_server:
-                    service_base_function_id = container_args.app_layout.function_ids[function_def.function_name]
-                    service_function_hydration_data = [
-                        o for o in container_args.app_layout.objects if o.object_id == service_base_function_id
-                    ][0]
-                    service = import_server_service(
-                        function_def,
-                        service_function_hydration_data,
-                        client,
-                        ser_usr_cls,
-                    )
-                elif function_def.is_class:
-                    # this is a bit ugly - match the function and class based on function name to get metadata
-                    # This metadata is required in order to hydrate the class in case it's not globally
-                    # decorated (or serialized)
-                    service_base_function_id = container_args.app_layout.function_ids[function_def.function_name]
-                    service_function_hydration_data = [
-                        o for o in container_args.app_layout.objects if o.object_id == service_base_function_id
-                    ][0]
-                    class_id = container_args.app_layout.class_ids[function_def.function_name.removesuffix(".*")]
-
-                    service = import_class_service(
-                        function_def,
-                        service_function_hydration_data,
-                        class_id,
-                        client,
-                        ser_usr_cls,
-                        param_args,
-                        param_kwargs,
-                    )
-                else:
-                    assert ser_usr_cls is None
-                    service = import_single_function_service(
-                        function_def,
-                        ser_fun,
-                    )
-
-            active_app = service.app
-
-            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 = running_app_from_layout(container_args.app_id, 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
-        app: App = cast(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 service.service_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.service_deps) != len(dep_object_ids):
-                raise ExecutionError(
-                    f"Function has {len(service.service_deps)} dependencies"
-                    f" but container got {len(dep_object_ids)} object ids.\n"
-                    f"Code deps: {service.service_deps}\n"
-                    f"Object ids: {dep_object_ids}\n"
-                    "\n"
-                    "This can happen if you are defining Modal objects under a conditional statement "
-                    "that evaluates differently in the local and remote environments."
-                )
-            for object_id, obj in zip(dep_object_ids, service.service_deps):
-                metadata: Message | None = container_app.object_handle_metadata[object_id]
-                obj._hydrate(object_id, _client, metadata)
+        service: Service = hydrate_function(
+            container_args, container_io_manager.get_task_lifecycle_manager(), function_def, _client, client
+        )
+        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
 
         # Initialize clustered functions.
         if function_def._experimental_group_size > 0:
@@ -407,10 +466,14 @@ def main(container_args: api_pb2.ContainerArguments, client: Client):
         ) as finalized_functions:
             # This context managers handles pre/post snapshot lifecycle of the user code,
             # finalized functions, ASGI/WSGI lifespan, and volume commit.
-            if function_def.is_server:
-                call_server(event_loop)
-            else:
-                call_function(event_loop, container_io_manager, finalized_functions, batch_max_size, batch_wait_ms)
+            call_function(event_loop, container_io_manager, finalized_functions, batch_max_size, batch_wait_ms)
+
+
+def main(container_args: api_pb2.ContainerArguments, client: Client):
+    if container_args.function_def.is_server:
+        run_server(container_args, client)
+    else:
+        run_function(container_args, client)
 
 
 if __name__ == "__main__":

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/_functions.py

@@ -19,6 +19,7 @@ from modal_proto import api_pb2
 from modal_proto.modal_api_grpc import ModalClientModal
 
 from ._function_variants import _FunctionOptions, _make_function_variant
+from ._image import _Image
 from ._load_context import LoadContext
 from ._object import _Object, live_method, live_method_gen
 from ._output.pty import get_pty_info
@@ -65,7 +66,6 @@ from .exception import (
     NotFoundError,
     OutputExpiredError,
 )
-from .image import _Image
 from .mount import _get_client_mount, _Mount
 from .network_file_system import _NetworkFileSystem, network_file_system_mount_protos
 from .output import OutputManager

modal-1.4.4.dev14/modal/image.py → modal-1.4.4.dev16/modal/_image.py

@@ -8,7 +8,7 @@ import sys
 import typing
 import warnings
 from collections.abc import Callable, Collection, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from inspect import isfunction
 from pathlib import Path, PurePosixPath
 from typing import (
@@ -32,7 +32,7 @@ from ._load_context import LoadContext
 from ._object import _Object, live_method_gen
 from ._resolver import Resolver
 from ._serialization import get_preferred_payload_format, serialize
-from ._utils.async_utils import TaskContext, deprecate_aio_usage, synchronize_api, synchronizer
+from ._utils.async_utils import TaskContext, deprecate_aio_usage, synchronizer
 from ._utils.blob_utils import MAX_OBJECT_SIZE_BYTES
 from ._utils.docker_utils import (
     extract_copy_command_patterns,
@@ -362,9 +362,8 @@ class _ImageRegistryConfig:
 
 @dataclass
 class DockerfileSpec:
-    # Ideally we would use field() with default_factory=, but doesn't work with synchronicity type-stub gen
-    commands: list[str]
-    context_files: dict[str, str]
+    commands: list[str] = field(default_factory=list)
+    context_files: dict[str, str] = field(default_factory=dict)
 
 
 async def _image_await_build_result(image_id: str, client: _Client) -> api_pb2.ImageJoinStreamingResponse:
@@ -538,7 +537,7 @@ class _Image(_Object, type_prefix="im"):
             raise InvalidError("Cannot run a build function with multiple base images!")
 
         def _deps() -> Sequence[_Object]:
-            deps = tuple(base_images.values()) + tuple(secrets)
+            deps: tuple[_Object, ...] = tuple(base_images.values()) + tuple(secrets)
             if build_function:
                 deps += (build_function,)
             if image_registry_config and image_registry_config.secret:
@@ -670,10 +669,10 @@ class _Image(_Object, type_prefix="im"):
             else:
                 # not built or in the process of building - wait for build
                 logger.debug("Waiting for image %s" % image_id)
-                resp = await _image_await_build_result(image_id, load_context.client)
-                result = resp.result
-                if resp.HasField("metadata"):
-                    metadata = resp.metadata
+                build_resp = await _image_await_build_result(image_id, load_context.client)
+                result = build_resp.result
+                if build_resp.HasField("metadata"):
+                    metadata = build_resp.metadata
 
             if result.status == api_pb2.GenericResult.GENERIC_STATUS_FAILURE:
                 if result.exception:
@@ -701,7 +700,7 @@ class _Image(_Object, type_prefix="im"):
                 raise RemoteError("Unknown status %s!" % result.status)
 
             self._hydrate(image_id, load_context.client, metadata)
-            local_mounts = set()
+            local_mounts: set[_Mount] = set()
             for base in base_images.values():
                 local_mounts |= base._serve_mounts
             if context_mount and context_mount.is_local():
@@ -1865,8 +1864,8 @@ class _Image(_Object, type_prefix="im"):
         """
         if not isinstance(entrypoint_commands, list) or not all(isinstance(x, str) for x in entrypoint_commands):
             raise InvalidError("entrypoint_commands must be a list of strings.")
-        args_str = _flatten_str_args("entrypoint", "entrypoint_commands", entrypoint_commands)
-        args_str = '"' + '", "'.join(args_str) + '"' if args_str else ""
+        args = _flatten_str_args("entrypoint", "entrypoint_commands", entrypoint_commands)
+        args_str = '"' + '", "'.join(args) + '"' if args else ""
         dockerfile_cmd = f"ENTRYPOINT [{args_str}]"
 
         return self.dockerfile_commands(dockerfile_cmd)
@@ -1885,8 +1884,8 @@ class _Image(_Object, type_prefix="im"):
         """
         if not isinstance(shell_commands, list) or not all(isinstance(x, str) for x in shell_commands):
             raise InvalidError("shell_commands must be a list of strings.")
-        args_str = _flatten_str_args("shell", "shell_commands", shell_commands)
-        args_str = '"' + '", "'.join(args_str) + '"' if args_str else ""
+        args = _flatten_str_args("shell", "shell_commands", shell_commands)
+        args_str = '"' + '", "'.join(args) + '"' if args else ""
         dockerfile_cmd = f"SHELL [{args_str}]"
 
         return self.dockerfile_commands(dockerfile_cmd)
@@ -2639,7 +2638,7 @@ class _Image(_Object, type_prefix="im"):
         function = _Function.from_local(
             info,
             app=None,
-            image=self,  # type: ignore[reportArgumentType]  # TODO: probably conflict with type stub?
+            image=self,
             secrets=secrets,
             gpu=gpu,
             volumes=volumes,
@@ -2756,8 +2755,8 @@ class _Image(_Object, type_prefix="im"):
         if not isinstance(cmd, list) or not all(isinstance(x, str) for x in cmd):
             raise InvalidError("Image CMD must be a list of strings.")
 
-        cmd_str = _flatten_str_args("cmd", "cmd", cmd)
-        cmd_str = '"' + '", "'.join(cmd_str) + '"' if cmd_str else ""
+        cmd_args = _flatten_str_args("cmd", "cmd", cmd)
+        cmd_str = '"' + '", "'.join(cmd_args) + '"' if cmd_args else ""
         dockerfile_cmd = f"CMD [{cmd_str}]"
 
         return self.dockerfile_commands(dockerfile_cmd)
@@ -2849,6 +2848,3 @@ class _Image(_Object, type_prefix="im"):
                 "Images cannot currently be hydrated on demand; you can build an Image by running an App that uses it."
             )
         return self
-
-
-Image = synchronize_api(_Image)

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/_runtime/execution_context.py

@@ -7,6 +7,7 @@ from modal._utils.async_utils import synchronize_api
 from modal.exception import InvalidError
 
 from .container_io_manager import _ContainerIOManager
+from .task_lifecycle_manager import _TaskLifecycleManager
 
 
 def is_local() -> bool:
@@ -18,7 +19,7 @@ def is_local() -> bool:
     even though those processes are running on Modal hardware.
 
     """
-    return not _ContainerIOManager._singleton
+    return _TaskLifecycleManager._singleton is None
 
 
 async def _interact() -> None:
@@ -29,7 +30,7 @@ async def _interact() -> None:
     """
     container_io_manager = _ContainerIOManager._singleton
     if not container_io_manager:
-        raise InvalidError("Interactivity only works inside a Modal container.")
+        raise InvalidError("Interactivity only works inside a Modal Function's container.")
     else:
         await container_io_manager.interact()
 

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/_runtime/task_lifecycle_manager.py

@@ -7,6 +7,10 @@ import sys
 import traceback
 from collections.abc import AsyncGenerator
 from pathlib import Path
+from typing import (
+    ClassVar,
+    Optional,
+)
 
 from synchronicity.async_wrap import asynccontextmanager
 
@@ -30,6 +34,8 @@ class UserException(Exception):
 
 
 class _TaskLifecycleManager:
+    _singleton: ClassVar[Optional["_TaskLifecycleManager"]] = None
+
     task_id: str
     function_id: str
     function_def: api_pb2.Function
@@ -37,7 +43,7 @@ class _TaskLifecycleManager:
     _client: _Client
     _cuda_checkpoint_session: gpu_memory_snapshot.CudaCheckpointSession | None
 
-    def __init__(
+    def _init(
         self,
         task_id: str,
         function_id: str,
@@ -52,6 +58,23 @@ class _TaskLifecycleManager:
         self._client = client
         self._cuda_checkpoint_session = None
 
+    def __new__(
+        cls,
+        task_id: str,
+        function_id: str,
+        function_def: api_pb2.Function,
+        checkpoint_id: str | None,
+        client: _Client,
+    ) -> "_TaskLifecycleManager":
+        cls._singleton = super().__new__(cls)
+        cls._singleton._init(task_id, function_id, function_def, checkpoint_id, client)
+        return cls._singleton
+
+    @classmethod
+    def _reset_singleton(cls) -> None:
+        """Only used for tests."""
+        cls._singleton = None
+
     @asynccontextmanager
     async def handle_task_lifecycle_exception(
         self,

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/_runtime/task_lifecycle_manager.pyi

@@ -11,6 +11,7 @@ class UserException(Exception):
     ...
 
 class _TaskLifecycleManager:
+    _singleton: typing.ClassVar[typing.Optional[_TaskLifecycleManager]]
     task_id: str
     function_id: str
     function_def: modal_proto.api_pb2.Function
@@ -18,15 +19,29 @@ class _TaskLifecycleManager:
     _client: modal.client._Client
     _cuda_checkpoint_session: typing.Optional[modal._runtime.gpu_memory_snapshot.CudaCheckpointSession]
 
-    def __init__(
+    def _init(
         self,
         task_id: str,
         function_id: str,
         function_def: modal_proto.api_pb2.Function,
         checkpoint_id: typing.Optional[str],
         client: modal.client._Client,
-    ) -> None:
-        """Initialize self.  See help(type(self)) for accurate signature."""
+    ) -> None: ...
+    @staticmethod
+    def __new__(
+        cls,
+        task_id: str,
+        function_id: str,
+        function_def: modal_proto.api_pb2.Function,
+        checkpoint_id: typing.Optional[str],
+        client: modal.client._Client,
+    ) -> _TaskLifecycleManager:
+        """Create and return a new object.  See help(type) for accurate signature."""
+        ...
+
+    @classmethod
+    def _reset_singleton(cls) -> None:
+        """Only used for tests."""
         ...
 
     def handle_task_lifecycle_exception(self) -> typing.AsyncContextManager[None]:
@@ -45,6 +60,7 @@ class _TaskLifecycleManager:
         ...
 
 class TaskLifecycleManager:
+    _singleton: typing.ClassVar[typing.Optional[TaskLifecycleManager]]
     task_id: str
     function_id: str
     function_def: modal_proto.api_pb2.Function
@@ -52,7 +68,11 @@ class TaskLifecycleManager:
     _client: modal.client.Client
     _cuda_checkpoint_session: typing.Optional[modal._runtime.gpu_memory_snapshot.CudaCheckpointSession]
 
-    def __init__(
+    def __init__(self, /, *args, **kwargs):
+        """Initialize self.  See help(type(self)) for accurate signature."""
+        ...
+
+    def _init(
         self,
         task_id: str,
         function_id: str,
@@ -60,6 +80,10 @@ class TaskLifecycleManager:
         checkpoint_id: typing.Optional[str],
         client: modal.client.Client,
     ) -> None: ...
+    @classmethod
+    def _reset_singleton(cls) -> None:
+        """Only used for tests."""
+        ...
 
     class __handle_task_lifecycle_exception_spec(typing_extensions.Protocol):
         def __call__(self, /) -> synchronicity.combined_types.AsyncAndBlockingContextManager[None]:

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/app.py

@@ -18,6 +18,7 @@ from synchronicity.async_wrap import asynccontextmanager
 from modal_proto import api_pb2
 
 from ._functions import _Function
+from ._image import _Image
 from ._ipython import is_interactive_ipython
 from ._load_context import LoadContext
 from ._object import _get_environment_name, _Object
@@ -47,7 +48,6 @@ from .cls import _Cls, parameter
 from .config import logger
 from .exception import ExecutionError, InvalidError
 from .functions import Function
-from .image import _Image
 from .network_file_system import _NetworkFileSystem
 from .partial_function import PartialFunction
 from .proxy import _Proxy

{modal-1.4.4.dev14 → modal-1.4.4.dev16}/modal/app.pyi

@@ -1,5 +1,6 @@
 import collections.abc
 import modal._functions
+import modal._image
 import modal._load_context
 import modal._partial_function
 import modal._server
@@ -80,7 +81,7 @@ class _LocalAppState:
 
     functions: dict[str, modal._functions._Function]
     classes: dict[str, modal.cls._Cls]
-    image_default: typing.Optional[modal.image._Image]
+    image_default: typing.Optional[modal._image._Image]
     web_endpoints: list[str]
     local_entrypoints: dict[str, _LocalEntrypoint]
     tags: dict[str, str]
@@ -92,7 +93,7 @@ class _LocalAppState:
         self,
         functions: dict[str, modal._functions._Function],
         classes: dict[str, modal.cls._Cls],
-        image_default: typing.Optional[modal.image._Image],
+        image_default: typing.Optional[modal._image._Image],
         web_endpoints: list[str],
         local_entrypoints: dict[str, _LocalEntrypoint],
         tags: dict[str, str],
@@ -163,7 +164,7 @@ class _App:
         name: typing.Optional[str] = None,
         *,
         tags: typing.Optional[dict[str, str]] = None,
-        image: typing.Optional[modal.image._Image] = None,
+        image: typing.Optional[modal._image._Image] = None,
         secrets: collections.abc.Sequence[modal.secret._Secret] = [],
         volumes: dict[typing.Union[str, pathlib.PurePosixPath], modal.volume._Volume] = {},
         include_source: bool = True,
@@ -282,7 +283,7 @@ class _App:
 
     def _validate_blueprint_value(self, key: str, value: typing.Any): ...
     @property
-    def image(self) -> modal.image._Image:
+    def image(self) -> modal._image._Image:
         """mdmd:hidden
         Retrieve the Image that will be used as the default for any Functions registered to the App.
 
@@ -546,7 +547,7 @@ class _App:
         self,
         _warn_parentheses_missing=None,
         *,
-        image: typing.Optional[modal.image._Image] = None,
+        image: typing.Optional[modal._image._Image] = None,
         schedule: typing.Optional[modal.schedule.Schedule] = None,
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret._Secret]] = None,
@@ -645,7 +646,7 @@ class _App:
         self,
         _warn_parentheses_missing=None,
         *,
-        image: typing.Optional[modal.image._Image] = None,
+        image: typing.Optional[modal._image._Image] = None,
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret._Secret]] = None,
         gpu: typing.Union[str, list[str], None] = None,
@@ -731,7 +732,7 @@ class _App:
         self,
         _warn_parentheses_missing=None,
         *,
-        image: typing.Optional[modal.image._Image] = None,
+        image: typing.Optional[modal._image._Image] = None,
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret._Secret]] = None,
         gpu: typing.Union[str, list[str], None] = None,
@@ -1948,4 +1949,4 @@ class App:
         """Only used for tests."""
         ...
 
-_default_image: modal.image._Image
+_default_image: modal._image._Image