modal 1.2.2.dev30__py3-none-any.whl → 1.2.2.dev31__py3-none-any.whl

modal/runner.py

@@ -8,7 +8,6 @@ import asyncio
 import dataclasses
 import os
 import time
-import typing
 from collections.abc import AsyncGenerator
 from contextlib import nullcontext
 from multiprocessing.synchronize import Event
@@ -19,6 +18,7 @@ from synchronicity.async_wrap import asynccontextmanager
 
 import modal._runtime.execution_context
 import modal_proto.api_pb2
+from modal._load_context import LoadContext
 from modal._utils.grpc_utils import Retry
 from modal_proto import api_pb2
 
@@ -125,18 +125,14 @@ async def _init_local_app_from_name(
 
 
 async def _create_all_objects(
-    client: _Client,
     running_app: RunningApp,
     local_app_state: "modal.app._LocalAppState",
-    environment_name: str,
+    load_context: LoadContext,
 ) -> None:
     """Create objects that have been defined but not created on the server."""
     indexed_objects: dict[str, _Object] = {**local_app_state.functions, **local_app_state.classes}
-    resolver = Resolver(
-        client,
-        environment_name=environment_name,
-        app_id=running_app.app_id,
-    )
+
+    resolver = Resolver()
     with resolver.display():
         # Get current objects, and reset all objects
         tag_to_object_id = {**running_app.function_ids, **running_app.class_ids}
@@ -159,7 +155,7 @@ async def _create_all_objects(
             # Note: preload only currently implemented for Functions, returns None otherwise
             # this is to ensure that directly referenced functions from the global scope has
             # ids associated with them when they are serialized into other functions
-            await resolver.preload(obj, existing_object_id)
+            await resolver.preload(obj, load_context, existing_object_id)
             if obj.is_hydrated:
                 tag_to_object_id[tag] = obj.object_id
 
@@ -167,7 +163,8 @@ async def _create_all_objects(
 
         async def _load(tag, obj):
             existing_object_id = tag_to_object_id.get(tag)
-            await resolver.load(obj, existing_object_id)
+            # Pass load_context so dependencies can inherit app_id, client, etc.
+            await resolver.load(obj, load_context, existing_object_id=existing_object_id)
             if _Function._is_id_type(obj.object_id):
                 running_app.function_ids[tag] = obj.object_id
             elif _Cls._is_id_type(obj.object_id):
@@ -266,8 +263,9 @@ async def _run_app(
     interactive: bool = False,
 ) -> AsyncGenerator["modal.app._App", None]:
     """mdmd:hidden"""
-    if environment_name is None:
-        environment_name = typing.cast(str, config.get("environment"))
+    load_context = await app._root_load_context.reset().in_place_upgrade(
+        client=client, environment_name=environment_name
+    )
 
     if modal._runtime.execution_context._is_currently_importing:
         raise InvalidError("Can not run an app in global scope within a container")
@@ -288,9 +286,6 @@ async def _run_app(
             # https://docs.python.org/3/library/__main__.html#import-main
             app.set_description(__main__.__name__)
 
-    if client is None:
-        client = await _Client.from_env()
-
     app_state = api_pb2.APP_STATE_DETACHED if detach else api_pb2.APP_STATE_EPHEMERAL
 
     output_mgr = _get_output_manager()
@@ -301,21 +296,22 @@ async def _run_app(
     local_app_state = app._local_state
 
     running_app: RunningApp = await _init_local_app_new(
-        client,
+        load_context.client,
         app.description or "",
         local_app_state.tags,
-        environment_name=environment_name or "",
+        environment_name=load_context.environment_name,
         app_state=app_state,
         interactive=interactive,
     )
+    await load_context.in_place_upgrade(app_id=running_app.app_id)
 
     logs_timeout = config["logs_timeout"]
-    async with app._set_local_app(client, running_app), TaskContext(grace=logs_timeout) as tc:
+    async with app._set_local_app(load_context.client, running_app), TaskContext(grace=logs_timeout) as tc:
         # Start heartbeats loop to keep the client alive
         # we don't log heartbeat exceptions in detached mode
         # as losing the local connection will not affect the running app
         def heartbeat():
-            return _heartbeat(client, running_app.app_id)
+            return _heartbeat(load_context.client, running_app.app_id)
 
         heartbeat_loop = tc.infinite_loop(heartbeat, sleep=HEARTBEAT_INTERVAL, log_exception=not detach)
         logs_loop: Optional[asyncio.Task] = None
@@ -336,24 +332,26 @@ async def _run_app(
             # Start logs loop
 
             logs_loop = tc.create_task(
-                get_app_logs_loop(client, output_mgr, app_id=running_app.app_id, app_logs_url=running_app.app_logs_url)
+                get_app_logs_loop(
+                    load_context.client, output_mgr, app_id=running_app.app_id, app_logs_url=running_app.app_logs_url
+                )
             )
 
         try:
             # Create all members
-            await _create_all_objects(client, running_app, local_app_state, environment_name)
+            await _create_all_objects(running_app, local_app_state, load_context)
 
             # Publish the app
-            await _publish_app(client, running_app, app_state, local_app_state)
+            await _publish_app(load_context.client, running_app, app_state, local_app_state)
         except asyncio.CancelledError as e:
             # this typically happens on sigint/ctrl-C during setup (the KeyboardInterrupt happens in the main thread)
             if output_mgr := _get_output_manager():
                 output_mgr.print("Aborting app initialization...\n")
 
-            await _status_based_disconnect(client, running_app.app_id, e)
+            await _status_based_disconnect(load_context.client, running_app.app_id, e)
             raise
         except BaseException as e:
-            await _status_based_disconnect(client, running_app.app_id, e)
+            await _status_based_disconnect(load_context.client, running_app.app_id, e)
             raise
 
         detached_disconnect_msg = (
@@ -381,7 +379,7 @@ async def _run_app(
                 yield app
             # successful completion!
             heartbeat_loop.cancel()
-            await _status_based_disconnect(client, running_app.app_id, exc_info=None)
+            await _status_based_disconnect(load_context.client, running_app.app_id, exc_info=None)
         except KeyboardInterrupt as e:
             # this happens only if sigint comes in during the yield block above
             if detach:
@@ -390,13 +388,13 @@ async def _run_app(
                     output_mgr.print(detached_disconnect_msg)
                 if logs_loop:
                     logs_loop.cancel()
-                await _status_based_disconnect(client, running_app.app_id, e)
+                await _status_based_disconnect(load_context.client, running_app.app_id, e)
             else:
                 if output_mgr := _get_output_manager():
                     output_mgr.print(
                         "Disconnecting from Modal - This will terminate your Modal app in a few seconds.\n"
                     )
-                await _status_based_disconnect(client, running_app.app_id, e)
+                await _status_based_disconnect(load_context.client, running_app.app_id, e)
                 if logs_loop:
                     try:
                         await asyncio.wait_for(logs_loop, timeout=logs_timeout)
@@ -421,7 +419,7 @@ async def _run_app(
             raise
         except BaseException as e:
             logger.info("Exception during app run")
-            await _status_based_disconnect(client, running_app.app_id, e)
+            await _status_based_disconnect(load_context.client, running_app.app_id, e)
             raise
 
         # wait for logs gracefully, even though the task context would do the same
@@ -449,21 +447,17 @@ async def _serve_update(
 ) -> None:
     """mdmd:hidden"""
     # Used by child process to reinitialize a served app
-    client = await _Client.from_env()
+    load_context = await app._root_load_context.reset().in_place_upgrade(environment_name=environment_name)
     try:
-        running_app: RunningApp = await _init_local_app_existing(client, existing_app_id, environment_name)
+        running_app: RunningApp = await _init_local_app_existing(load_context.client, existing_app_id, environment_name)
+        await load_context.in_place_upgrade(app_id=running_app.app_id)
         local_app_state = app._local_state
         # Create objects
-        await _create_all_objects(
-            client,
-            running_app,
-            local_app_state,
-            environment_name,
-        )
+        await _create_all_objects(running_app, local_app_state, load_context)
 
         # Publish the updated app
         await _publish_app(
-            client,
+            load_context.client,
             running_app,
             app_state=api_pb2.APP_STATE_UNSPECIFIED,
             app_local_state=local_app_state,
@@ -498,9 +492,6 @@ async def _deploy_app(
 
     Users should prefer the `modal deploy` CLI or the `App.deploy` method.
     """
-    if environment_name is None:
-        environment_name = typing.cast(str, config.get("environment"))
-
     warn_if_passing_namespace(namespace, "modal.runner.deploy_app")
 
     name = name or app.name or ""
@@ -532,8 +523,18 @@ async def _deploy_app(
     # Get git information to track deployment history
     commit_info_task = asyncio.create_task(get_git_commit_info())
 
+    # We need to do in-place replacement of fields in self._root_load_context in case it has already "spread"
+    # to with_options() instances or similar before load
+    root_load_context = await app._root_load_context.reset().in_place_upgrade(
+        client=client,
+        environment_name=environment_name,
+    )
     running_app: RunningApp = await _init_local_app_from_name(
-        client, name, local_app_state.tags, environment_name=environment_name
+        root_load_context.client, name, local_app_state.tags, environment_name=root_load_context.environment_name
+    )
+
+    await root_load_context.in_place_upgrade(
+        app_id=running_app.app_id,
     )
 
     async with TaskContext(0) as tc:
@@ -545,12 +546,7 @@ async def _deploy_app(
 
         try:
             # Create all members
-            await _create_all_objects(
-                client,
-                running_app,
-                local_app_state,
-                environment_name=environment_name,
-            )
+            await _create_all_objects(running_app, local_app_state, root_load_context)
 
             commit_info = None
             try:

modal/runner.pyi

@@ -1,3 +1,4 @@
+import modal._load_context
 import modal.app
 import modal.client
 import modal.running_app
@@ -25,10 +26,9 @@ async def _init_local_app_from_name(
     client: modal.client._Client, name: str, tags: dict[str, str], environment_name: str = ""
 ) -> modal.running_app.RunningApp: ...
 async def _create_all_objects(
-    client: modal.client._Client,
     running_app: modal.running_app.RunningApp,
     local_app_state: modal.app._LocalAppState,
-    environment_name: str,
+    load_context: modal._load_context.LoadContext,
 ) -> None:
     """Create objects that have been defined but not created on the server."""
     ...

modal/sandbox.py

@@ -23,6 +23,7 @@ from modal.mount import _Mount
 from modal.volume import _Volume
 from modal_proto import api_pb2, task_command_router_pb2 as sr_pb2
 
+from ._load_context import LoadContext
 from ._object import _get_environment_name, _Object
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
@@ -191,7 +192,9 @@ class _Sandbox(_Object, type_prefix="sb"):
                 deps.append(proxy)
             return deps
 
-        async def _load(self: _Sandbox, resolver: Resolver, _existing_object_id: Optional[str]):
+        async def _load(
+            self: _Sandbox, resolver: Resolver, load_context: LoadContext, _existing_object_id: Optional[str]
+        ):
             # Relies on dicts being ordered (true as of Python 3.6).
             volume_mounts = [
                 api_pb2.VolumeMount(
@@ -260,18 +263,18 @@ class _Sandbox(_Object, type_prefix="sb"):
                 experimental_options=experimental_options,
             )
 
-            create_req = api_pb2.SandboxCreateRequest(app_id=resolver.app_id, definition=definition)
+            create_req = api_pb2.SandboxCreateRequest(app_id=load_context.app_id, definition=definition)
             try:
-                create_resp = await resolver.client.stub.SandboxCreate(create_req)
+                create_resp = await load_context.client.stub.SandboxCreate(create_req)
             except GRPCError as exc:
                 if exc.status == Status.ALREADY_EXISTS:
                     raise AlreadyExistsError(exc.message)
                 raise exc
 
             sandbox_id = create_resp.sandbox_id
-            self._hydrate(sandbox_id, resolver.client, None)
+            self._hydrate(sandbox_id, load_context.client, None)
 
-        return _Sandbox._from_loader(_load, "Sandbox()", deps=_deps)
+        return _Sandbox._from_loader(_load, "Sandbox()", deps=_deps, load_context_overrides=LoadContext.empty())
 
     @staticmethod
     async def create(
@@ -486,6 +489,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             app_id = app.app_id
             app_client = app._client
         elif (container_app := _App._get_container_app()) is not None:
+            # implicit app/client provided by running in a modal Function
             app_id = container_app.app_id
             app_client = container_app._client
         else:
@@ -498,10 +502,11 @@ class _Sandbox(_Object, type_prefix="sb"):
                 "```",
             )
 
-        client = client or app_client or await _Client.from_env()
+        client = client or app_client
 
-        resolver = Resolver(client, app_id=app_id)
-        await resolver.load(obj)
+        resolver = Resolver()
+        load_context = LoadContext(client=client, app_id=app_id)
+        await resolver.load(obj, load_context)
         return obj
 
     def _hydrate_metadata(self, handle_metadata: Optional[Message]):
@@ -606,12 +611,13 @@ class _Sandbox(_Object, type_prefix="sb"):
         image_id = resp.image_id
         metadata = resp.image_metadata
 
-        async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: _Image, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             # no need to hydrate again since we do it eagerly below
             pass
 
         rep = "Image()"
-        image = _Image._from_loader(_load, rep, hydrate_lazily=True)
+        # TODO: use ._new_hydrated instead
+        image = _Image._from_loader(_load, rep, hydrate_lazily=True, load_context_overrides=LoadContext.empty())
         image._hydrate(image_id, self._client, metadata)  # hydrating eagerly since we have all of the data
 
         return image
@@ -990,12 +996,15 @@ class _Sandbox(_Object, type_prefix="sb"):
         if wait_resp.result.status != api_pb2.GenericResult.GENERIC_STATUS_SUCCESS:
             raise ExecutionError(wait_resp.result.exception)
 
-        async def _load(self: _SandboxSnapshot, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _SandboxSnapshot, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             # we eagerly hydrate the sandbox snapshot below
             pass
 
         rep = "SandboxSnapshot()"
-        obj = _SandboxSnapshot._from_loader(_load, rep, hydrate_lazily=True)
+        # TODO: use ._new_hydrated instead
+        obj = _SandboxSnapshot._from_loader(_load, rep, hydrate_lazily=True, load_context_overrides=LoadContext.empty())
         obj._hydrate(snapshot_id, self._client, None)
 
         return obj

modal/secret.py

@@ -10,6 +10,7 @@ from synchronicity import classproperty
 
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import _get_environment_name, _Object, live_method
 from ._resolver import Resolver
 from ._runtime.execution_context import is_local
@@ -259,8 +260,10 @@ class _Secret(_Object, type_prefix="st"):
         if not all(isinstance(v, str) for v in env_dict_filtered.values()):
             raise InvalidError(ENV_DICT_WRONG_TYPE_ERR)
 
-        async def _load(self: _Secret, resolver: Resolver, existing_object_id: Optional[str]):
-            if resolver.app_id is not None:
+        async def _load(
+            self: _Secret, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
+            if load_context.app_id is not None:
                 object_creation_type = api_pb2.OBJECT_CREATION_TYPE_ANONYMOUS_OWNED_BY_APP
             else:
                 object_creation_type = api_pb2.OBJECT_CREATION_TYPE_EPHEMERAL
@@ -268,21 +271,22 @@ class _Secret(_Object, type_prefix="st"):
             req = api_pb2.SecretGetOrCreateRequest(
                 object_creation_type=object_creation_type,
                 env_dict=env_dict_filtered,
-                app_id=resolver.app_id,
-                environment_name=resolver.environment_name,
+                app_id=load_context.app_id,
+                environment_name=load_context.environment_name,
             )
             try:
-                resp = await resolver.client.stub.SecretGetOrCreate(req)
+                resp = await load_context.client.stub.SecretGetOrCreate(req)
             except GRPCError as exc:
                 if exc.status == Status.INVALID_ARGUMENT:
                     raise InvalidError(exc.message)
                 if exc.status == Status.FAILED_PRECONDITION:
                     raise InvalidError(exc.message)
                 raise
-            self._hydrate(resp.secret_id, resolver.client, resp.metadata)
+            self._hydrate(resp.secret_id, load_context.client, resp.metadata)
 
         rep = f"Secret.from_dict([{', '.join(env_dict.keys())}])"
-        return _Secret._from_loader(_load, rep, hydrate_lazily=True)
+        # TODO: scoping - these should probably not be lazily hydrated without having an app and/or sandbox association
+        return _Secret._from_loader(_load, rep, hydrate_lazily=True, load_context_overrides=LoadContext.empty())
 
     @staticmethod
     def from_local_environ(
@@ -330,7 +334,9 @@ class _Secret(_Object, type_prefix="st"):
         ```
         """
 
-        async def _load(self: _Secret, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _Secret, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             try:
                 from dotenv import dotenv_values, find_dotenv
                 from dotenv.main import _walk_to_root
@@ -359,13 +365,15 @@ class _Secret(_Object, type_prefix="st"):
             req = api_pb2.SecretGetOrCreateRequest(
                 object_creation_type=api_pb2.OBJECT_CREATION_TYPE_ANONYMOUS_OWNED_BY_APP,
                 env_dict=env_dict,
-                app_id=resolver.app_id,
+                app_id=load_context.app_id,  # TODO: what if app_id isn't set here (e.g. .hydrate())
             )
-            resp = await resolver.client.stub.SecretGetOrCreate(req)
+            resp = await load_context.client.stub.SecretGetOrCreate(req)
 
-            self._hydrate(resp.secret_id, resolver.client, resp.metadata)
+            self._hydrate(resp.secret_id, load_context.client, resp.metadata)
 
-        return _Secret._from_loader(_load, "Secret.from_dotenv()", hydrate_lazily=True)
+        return _Secret._from_loader(
+            _load, "Secret.from_dotenv()", hydrate_lazily=True, load_context_overrides=LoadContext.empty()
+        )
 
     @staticmethod
     def from_name(
@@ -376,6 +384,7 @@ class _Secret(_Object, type_prefix="st"):
         required_keys: list[
             str
         ] = [],  # Optionally, a list of required environment variables (will be asserted server-side)
+        client: Optional[_Client] = None,
     ) -> "_Secret":
         """Reference a Secret by its name.
 
@@ -393,23 +402,31 @@ class _Secret(_Object, type_prefix="st"):
         """
         warn_if_passing_namespace(namespace, "modal.Secret.from_name")
 
-        async def _load(self: _Secret, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _Secret, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             req = api_pb2.SecretGetOrCreateRequest(
                 deployment_name=name,
-                environment_name=_get_environment_name(environment_name, resolver),
+                environment_name=load_context.environment_name,
                 required_keys=required_keys,
             )
             try:
-                response = await resolver.client.stub.SecretGetOrCreate(req)
+                response = await load_context.client.stub.SecretGetOrCreate(req)
             except GRPCError as exc:
                 if exc.status == Status.NOT_FOUND:
                     raise NotFoundError(exc.message)
                 else:
                     raise
-            self._hydrate(response.secret_id, resolver.client, response.metadata)
+            self._hydrate(response.secret_id, load_context.client, response.metadata)
 
         rep = _Secret._repr(name, environment_name)
-        return _Secret._from_loader(_load, rep, hydrate_lazily=True, name=name)
+        return _Secret._from_loader(
+            _load,
+            rep,
+            hydrate_lazily=True,
+            name=name,
+            load_context_overrides=LoadContext(environment_name=environment_name, client=client),
+        )
 
     @staticmethod
     async def create_deployed(

modal/secret.pyi

@@ -423,7 +423,12 @@ class _Secret(modal._object._Object):
 
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, required_keys: list[str] = []
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        required_keys: list[str] = [],
+        client: typing.Optional[modal.client._Client] = None,
     ) -> _Secret:
         """Reference a Secret by its name.
 
@@ -543,7 +548,12 @@ class Secret(modal.object.Object):
 
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, required_keys: list[str] = []
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        required_keys: list[str] = [],
+        client: typing.Optional[modal.client.Client] = None,
     ) -> Secret:
         """Reference a Secret by its name.
 

modal/serving.py

@@ -4,13 +4,13 @@ import platform
 from collections.abc import AsyncGenerator
 from multiprocessing.context import SpawnProcess
 from multiprocessing.synchronize import Event
-from typing import TYPE_CHECKING, Optional, TypeVar
+from typing import TYPE_CHECKING, Optional
 
 from synchronicity.async_wrap import asynccontextmanager
 
 from modal._output import OutputManager
 
-from ._utils.async_utils import TaskContext, asyncify, synchronize_api, synchronizer
+from ._utils.async_utils import TaskContext, asyncify, synchronize_api
 from ._utils.logger import logger
 from ._watcher import watch
 from .cli.import_refs import ImportRef, import_app_from_ref
@@ -20,20 +20,16 @@ from .output import _get_output_manager, enable_output
 from .runner import _run_app, serve_update
 
 if TYPE_CHECKING:
-    from .app import _App
-else:
-    _App = TypeVar("_App")
+    import modal.app
 
 
 def _run_serve(
     import_ref: ImportRef, existing_app_id: str, is_ready: Event, environment_name: str, show_progress: bool
 ):
-    # subprocess entrypoint
-    _app = import_app_from_ref(import_ref, base_cmd="modal serve")
-    blocking_app = synchronizer._translate_out(_app)
+    app = import_app_from_ref(import_ref, base_cmd="modal serve")
 
     with enable_output(show_progress=show_progress):
-        serve_update(blocking_app, existing_app_id, is_ready, environment_name)
+        serve_update(app, existing_app_id, is_ready, environment_name)
 
 
 async def _restart_serve(
@@ -97,12 +93,12 @@ async def _run_watch_loop(
 
 @asynccontextmanager
 async def _serve_app(
-    app: "_App",
+    app: "modal.app._App",
     import_ref: ImportRef,
     *,
     _watcher: Optional[AsyncGenerator[set[str], None]] = None,  # for testing
     environment_name: Optional[str] = None,
-) -> AsyncGenerator["_App", None]:
+) -> AsyncGenerator["modal.app._App", None]:
     if environment_name is None:
         environment_name = config.get("environment")
 

modal/serving.pyi

@@ -1,4 +1,5 @@
 import collections.abc
+import modal.app
 import modal.cli.import_refs
 import multiprocessing.context
 import multiprocessing.synchronize
@@ -6,8 +7,6 @@ import synchronicity.combined_types
 import typing
 import typing_extensions
 
-_App = typing.TypeVar("_App")
-
 def _run_serve(
     import_ref: modal.cli.import_refs.ImportRef,
     existing_app_id: str,
@@ -27,31 +26,31 @@ async def _run_watch_loop(
     environment_name: str,
 ): ...
 def _serve_app(
-    app: _App,
+    app: modal.app._App,
     import_ref: modal.cli.import_refs.ImportRef,
     *,
     _watcher: typing.Optional[collections.abc.AsyncGenerator[set[str], None]] = None,
     environment_name: typing.Optional[str] = None,
-) -> typing.AsyncContextManager[_App]: ...
+) -> typing.AsyncContextManager[modal.app._App]: ...
 
 class __serve_app_spec(typing_extensions.Protocol):
     def __call__(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         import_ref: modal.cli.import_refs.ImportRef,
         *,
         _watcher: typing.Optional[typing.Generator[set[str], None, None]] = None,
         environment_name: typing.Optional[str] = None,
-    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[_App]: ...
+    ) -> synchronicity.combined_types.AsyncAndBlockingContextManager[modal.app.App]: ...
     def aio(
         self,
         /,
-        app: _App,
+        app: modal.app.App,
         import_ref: modal.cli.import_refs.ImportRef,
         *,
         _watcher: typing.Optional[collections.abc.AsyncGenerator[set[str], None]] = None,
         environment_name: typing.Optional[str] = None,
-    ) -> typing.AsyncContextManager[_App]: ...
+    ) -> typing.AsyncContextManager[modal.app.App]: ...
 
 serve_app: __serve_app_spec

modal/snapshot.py

@@ -3,6 +3,7 @@ from typing import Optional
 
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import _Object
 from ._resolver import Resolver
 from ._utils.async_utils import synchronize_api
@@ -23,14 +24,19 @@ class _SandboxSnapshot(_Object, type_prefix="sn"):
         """
         Construct a `SandboxSnapshot` object from a sandbox snapshot ID.
         """
-        if client is None:
-            client = await _Client.from_env()
+        # TODO: remove this - from_id constructor should not do io:
+        client = client or await _Client.from_env()
 
-        async def _load(self: _SandboxSnapshot, resolver: Resolver, existing_object_id: Optional[str]):
-            await client.stub.SandboxSnapshotGet(api_pb2.SandboxSnapshotGetRequest(snapshot_id=sandbox_snapshot_id))
+        async def _load(
+            self: _SandboxSnapshot, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
+            await load_context.client.stub.SandboxSnapshotGet(
+                api_pb2.SandboxSnapshotGetRequest(snapshot_id=sandbox_snapshot_id)
+            )
 
         rep = "SandboxSnapshot()"
-        obj = _SandboxSnapshot._from_loader(_load, rep)
+        obj = _SandboxSnapshot._from_loader(_load, rep, load_context_overrides=LoadContext(client=client))
+        # TODO: should this be a _Object._new_hydrated instead?
         obj._hydrate(sandbox_snapshot_id, client, None)
 
         return obj