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

modal/cli/dict.py

@@ -4,6 +4,7 @@ from typing import Optional
 import typer
 from typer import Argument, Option, Typer
 
+from modal._load_context import LoadContext
 from modal._output import make_console
 from modal._resolver import Resolver
 from modal._utils.async_utils import synchronizer
@@ -29,8 +30,10 @@ async def create(name: str, *, env: Optional[str] = ENV_OPTION):
     """
     d = _Dict.from_name(name, environment_name=env, create_if_missing=True)
     client = await _Client.from_env()
-    resolver = Resolver(client=client)
-    await resolver.load(d)
+    resolver = Resolver()
+
+    load_context = LoadContext(client=client, environment_name=env)
+    await resolver.load(d, load_context)
 
 
 @dict_cli.command(name="list", rich_help_panel="Management")

modal/cli/queues.py

@@ -5,6 +5,7 @@ from typing import Optional
 import typer
 from typer import Argument, Option, Typer
 
+from modal._load_context import LoadContext
 from modal._output import make_console
 from modal._resolver import Resolver
 from modal._utils.async_utils import synchronizer
@@ -38,8 +39,9 @@ async def create(name: str, *, env: Optional[str] = ENV_OPTION):
     """
     q = _Queue.from_name(name, environment_name=env, create_if_missing=True)
     client = await _Client.from_env()
-    resolver = Resolver(client=client)
-    await resolver.load(q)
+    resolver = Resolver()
+    load_context = LoadContext(client=client, environment_name=env)
+    await resolver.load(q, load_context)
 
 
 @queue_cli.command(name="delete", rich_help_panel="Management")

modal/client.pyi

@@ -33,7 +33,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.2.dev30",
+        version: str = "1.2.2.dev31",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -164,7 +164,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.2.2.dev30",
+        version: str = "1.2.2.dev31",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

modal/cls.py

@@ -12,6 +12,7 @@ from grpclib import GRPCError, Status
 from modal_proto import api_pb2
 
 from ._functions import _Function, _parse_retries
+from ._load_context import LoadContext
 from ._object import _Object, live_method
 from ._partial_function import (
     _find_callables_for_obj,
@@ -31,8 +32,8 @@ from ._utils.deprecation import (
     warn_on_renamed_autoscaler_settings,
 )
 from ._utils.mount_utils import validate_volumes
+from .client import _Client
 from .cloud_bucket_mount import _CloudBucketMount
-from .config import config
 from .exception import ExecutionError, InvalidError, NotFoundError
 from .gpu import GPU_T
 from .retries import Retries
@@ -134,7 +135,7 @@ def _bind_instance_method(cls: "_Cls", service_function: _Function, method_name:
         method_metadata = cls._method_metadata[method_name]
         new_function._hydrate(service_function.object_id, service_function.client, method_metadata)
 
-    async def _load(fun: "_Function", resolver: Resolver, existing_object_id: Optional[str]):
+    async def _load(fun: "_Function", resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
         # there is currently no actual loading logic executed to create each method on
         # the *parametrized* instance of a class - it uses the parameter-bound service-function
         # for the instance. This load method just makes sure to set all attributes after the
@@ -153,11 +154,14 @@ def _bind_instance_method(cls: "_Cls", service_function: _Function, method_name:
 
     rep = f"Method({cls._name}.{method_name})"
 
+    # Bound methods should *reference* their parent Cls's LoadContext
+    # so that it can be modified in place on the parent and be reflected in the method
     fun = _Function._from_loader(
         _load,
         rep,
         deps=_deps,
         hydrate_lazily=True,
+        load_context_overrides=cls._load_context_overrides,
     )
     if service_function.is_hydrated:
         # Eager hydration (skip load) if the instance service function is already loaded
@@ -421,14 +425,13 @@ class _Obj:
 
         # Not hydrated Cls, and we don't have the class - typically a Cls.from_name that
         # has not yet been loaded. So use a special loader that loads it lazily:
-        async def method_loader(fun, resolver: Resolver, existing_object_id):
-            await resolver.load(self._cls)  # load class so we get info about methods
+        async def method_loader(fun, resolver: Resolver, load_context: LoadContext, existing_object_id):
             method_function = _get_maybe_method()
             if method_function is None:
                 raise NotFoundError(
                     f"Class has no method {k}, and attributes can't be accessed for `Cls.from_name` instances"
                 )
-            await resolver.load(method_function)  # get the appropriate method handle (lazy)
+            await resolver.load(method_function, load_context)  # get the appropriate method handle (lazy)
             fun._hydrate_from_other(method_function)
 
         # The reason we don't *always* use this lazy loader is because it precludes attribute access
@@ -436,8 +439,9 @@ class _Obj:
         return _Function._from_loader(
             method_loader,
             rep=f"Method({self._cls._name}.{k})",
-            deps=lambda: [],  # TODO: use cls as dep instead of loading inside method_loader?
+            deps=lambda: [self._cls],
             hydrate_lazily=True,
+            load_context_overrides=self._cls._load_context_overrides,
         )
 
 
@@ -484,6 +488,7 @@ class _Cls(_Object, type_prefix="cs"):
         self._callables = other._callables
         self._name = other._name
         self._method_metadata = other._method_metadata
+        self._load_context_overrides = other._load_context_overrides
 
     def _get_partial_functions(self) -> dict[str, _PartialFunction]:
         if not self._user_cls:
@@ -595,15 +600,18 @@ More information on class parameterization can be found here: https://modal.com/
         def _deps() -> list[_Function]:
             return [class_service_function]
 
-        async def _load(self: "_Cls", resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: "_Cls", resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             req = api_pb2.ClassCreateRequest(
-                app_id=resolver.app_id, existing_class_id=existing_object_id, only_class_function=True
+                app_id=load_context.app_id, existing_class_id=existing_object_id, only_class_function=True
             )
-            resp = await resolver.client.stub.ClassCreate(req)
-            self._hydrate(resp.class_id, resolver.client, resp.handle_metadata)
+            resp = await load_context.client.stub.ClassCreate(req)
+            self._hydrate(resp.class_id, load_context.client, resp.handle_metadata)
 
         rep = f"Cls({user_cls.__name__})"
-        cls: _Cls = _Cls._from_loader(_load, rep, deps=_deps)
+        # Pass a *reference* to the App's LoadContext - this is important since the App is
+        # the only way to infer a LoadContext for an `@app.cls`, and the App doesn't
+        # get its client until *after* the Cls is created.
+        cls: _Cls = _Cls._from_loader(_load, rep, deps=_deps, load_context_overrides=app._root_load_context)
         cls._app = app
         cls._user_cls = user_cls
         cls._class_service_function = class_service_function
@@ -620,6 +628,7 @@ More information on class parameterization can be found here: https://modal.com/
         *,
         namespace: Any = None,  # mdmd:line-hidden
         environment_name: Optional[str] = None,
+        client: Optional["_Client"] = None,
     ) -> "_Cls":
         """Reference a Cls from a deployed App by its name.
 
@@ -632,19 +641,22 @@ More information on class parameterization can be found here: https://modal.com/
         ```
         """
         warn_if_passing_namespace(namespace, "modal.Cls.from_name")
-        _environment_name = environment_name or config.get("environment")
 
-        async def _load_remote(self: _Cls, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load_remote(
+            self: _Cls, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             request = api_pb2.ClassGetRequest(
                 app_name=app_name,
                 object_tag=name,
-                environment_name=_environment_name,
+                environment_name=load_context.environment_name,
                 only_class_function=True,
             )
             try:
-                response = await resolver.client.stub.ClassGet(request)
+                response = await load_context.client.stub.ClassGet(request)
             except NotFoundError as exc:
-                env_context = f" (in the '{environment_name}' environment)" if environment_name else ""
+                env_context = (
+                    f" (in the '{load_context.environment_name}' environment)" if load_context.environment_name else ""
+                )
                 raise NotFoundError(
                     f"Lookup failed for Cls '{name}' from the '{app_name}' app{env_context}: {exc}."
                 ) from None
@@ -655,19 +667,26 @@ More information on class parameterization can be found here: https://modal.com/
                     raise
 
             print_server_warnings(response.server_warnings)
-            await resolver.load(self._class_service_function)
-            self._hydrate(response.class_id, resolver.client, response.handle_metadata)
+            await resolver.load(self._class_service_function, load_context)
+            self._hydrate(response.class_id, load_context.client, response.handle_metadata)
 
         environment_rep = f", environment_name={environment_name!r}" if environment_name else ""
         rep = f"Cls.from_name({app_name!r}, {name!r}{environment_rep})"
-        cls = cls._from_loader(_load_remote, rep, is_another_app=True, hydrate_lazily=True)
+
+        load_context_overrides = LoadContext(client=client, environment_name=environment_name)
+        cls = cls._from_loader(
+            _load_remote,
+            rep,
+            is_another_app=True,
+            hydrate_lazily=True,
+            load_context_overrides=load_context_overrides,
+        )
 
         class_service_name = f"{name}.*"  # special name of the base service function for the class
         cls._class_service_function = _Function._from_name(
             app_name,
             class_service_name,
-            namespace=namespace,
-            environment_name=_environment_name,
+            load_context_overrides=load_context_overrides,
         )
         cls._name = name
         return cls
@@ -736,7 +755,7 @@ More information on class parameterization can be found here: https://modal.com/
                 " please use the `.with_concurrency` method instead.",
             )
 
-        async def _load_from_base(new_cls, resolver, existing_object_id):
+        async def _load_from_base(new_cls, resolver, load_context, existing_object_id):
             # this is a bit confusing, the cls will always have the same metadata
             # since it has the same *class* service function (i.e. "template")
             # But the (instance) service function for each Obj will be different
@@ -745,14 +764,21 @@ More information on class parameterization can be found here: https://modal.com/
             if not self.is_hydrated:
                 # this should only happen for Cls.from_name instances
                 # other classes should already be hydrated!
-                await resolver.load(self)
+                await resolver.load(self, load_context)
 
             new_cls._initialize_from_other(self)
 
         def _deps():
             return []
 
-        cls = _Cls._from_loader(_load_from_base, rep=f"{self._name}.with_options(...)", is_another_app=True, deps=_deps)
+        cls = _Cls._from_loader(
+            _load_from_base,
+            rep=f"{self._name}.with_options(...)",
+            is_another_app=True,
+            deps=_deps,
+            load_context_overrides=self._load_context_overrides,
+            hydrate_lazily=True,
+        )
         cls._initialize_from_other(self)
 
         # Validate volumes
@@ -797,16 +823,21 @@ More information on class parameterization can be found here: https://modal.com/
         ```
         """
 
-        async def _load_from_base(new_cls, resolver, existing_object_id):
+        async def _load_from_base(new_cls, resolver, load_context, existing_object_id):
             if not self.is_hydrated:
-                await resolver.load(self)
+                await resolver.load(self, load_context)
             new_cls._initialize_from_other(self)
 
         def _deps():
             return []
 
         cls = _Cls._from_loader(
-            _load_from_base, rep=f"{self._name}.with_concurrency(...)", is_another_app=True, deps=_deps
+            _load_from_base,
+            rep=f"{self._name}.with_concurrency(...)",
+            is_another_app=True,
+            deps=_deps,
+            load_context_overrides=self._load_context_overrides,
+            hydrate_lazily=True,
         )
         cls._initialize_from_other(self)
 
@@ -826,16 +857,21 @@ More information on class parameterization can be found here: https://modal.com/
         ```
         """
 
-        async def _load_from_base(new_cls, resolver, existing_object_id):
+        async def _load_from_base(new_cls, resolver, load_context, existing_object_id):
             if not self.is_hydrated:
-                await resolver.load(self)
+                await resolver.load(self, load_context)
             new_cls._initialize_from_other(self)
 
         def _deps():
             return []
 
         cls = _Cls._from_loader(
-            _load_from_base, rep=f"{self._name}.with_concurrency(...)", is_another_app=True, deps=_deps
+            _load_from_base,
+            rep=f"{self._name}.with_concurrency(...)",
+            is_another_app=True,
+            deps=_deps,
+            load_context_overrides=self._load_context_overrides,
+            hydrate_lazily=True,
         )
         cls._initialize_from_other(self)
 
@@ -865,17 +901,20 @@ More information on class parameterization can be found here: https://modal.com/
         # We create a synthetic dummy Function that is guaranteed to raise an AttributeError when
         # a user tries to use any of its "live methods" - this lets us raise exceptions for users
         # only if they try to access methods on a Cls as if they were methods on the instance.
-        async def method_loader(fun: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+        async def error_loader(
+            fun: _Function, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             raise AttributeError(
                 "You can't access methods on a Cls directly - Did you forget to instantiate the class first?\n"
                 "e.g. instead of MyClass.method.remote(), do MyClass().method.remote()"
             )
 
         return _Function._from_loader(
-            method_loader,
+            error_loader,
             rep=f"UnboundMethod({self._name}.{k})",
             deps=lambda: [],
             hydrate_lazily=True,
+            load_context_overrides=self._load_context_overrides,
         )
 
     def _is_local(self) -> bool:

modal/cls.pyi

@@ -5,6 +5,7 @@ import modal._functions
 import modal._object
 import modal._partial_function
 import modal.app
+import modal.client
 import modal.cloud_bucket_mount
 import modal.functions
 import modal.gpu
@@ -379,6 +380,7 @@ class _Cls(modal._object._Object):
         *,
         namespace: typing.Any = None,
         environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client._Client] = None,
     ) -> _Cls:
         """Reference a Cls from a deployed App by its name.
 
@@ -537,6 +539,7 @@ class Cls(modal.object.Object):
         *,
         namespace: typing.Any = None,
         environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> Cls:
         """Reference a Cls from a deployed App by its name.
 

modal/dict.py

@@ -12,6 +12,7 @@ from synchronicity.async_wrap import asynccontextmanager
 from modal._utils.grpc_utils import Retry
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import (
     EPHEMERAL_OBJECT_HEARTBEAT_SLEEP,
     _get_environment_name,
@@ -347,6 +348,7 @@ class _Dict(_Object, type_prefix="di"):
         namespace=None,  # mdmd:line-hidden
         environment_name: Optional[str] = None,
         create_if_missing: bool = False,
+        client: Optional[_Client] = None,
     ) -> "_Dict":
         """Reference a named Dict, creating if necessary.
 
@@ -368,20 +370,27 @@ class _Dict(_Object, type_prefix="di"):
                 "Passing data to `modal.Dict.from_name` is deprecated and will stop working in a future release.",
             )
 
-        async def _load(self: _Dict, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: _Dict, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             serialized = _serialize_dict(data if data is not None else {})
             req = api_pb2.DictGetOrCreateRequest(
                 deployment_name=name,
-                environment_name=_get_environment_name(environment_name, resolver),
+                environment_name=load_context.environment_name,
                 object_creation_type=(api_pb2.OBJECT_CREATION_TYPE_CREATE_IF_MISSING if create_if_missing else None),
                 data=serialized,
             )
-            response = await resolver.client.stub.DictGetOrCreate(req)
+            response = await load_context.client.stub.DictGetOrCreate(req)
             logger.debug(f"Created dict with id {response.dict_id}")
-            self._hydrate(response.dict_id, resolver.client, response.metadata)
+            self._hydrate(response.dict_id, load_context.client, response.metadata)
 
         rep = _Dict._repr(name, environment_name)
-        return _Dict._from_loader(_load, rep, is_another_app=True, hydrate_lazily=True, name=name)
+        return _Dict._from_loader(
+            _load,
+            rep,
+            is_another_app=True,
+            hydrate_lazily=True,
+            name=name,
+            load_context_overrides=LoadContext(environment_name=environment_name, client=client),
+        )
 
     @staticmethod
     async def delete(

modal/dict.pyi

@@ -448,6 +448,7 @@ class _Dict(modal._object._Object):
         namespace=None,
         environment_name: typing.Optional[str] = None,
         create_if_missing: bool = False,
+        client: typing.Optional[modal.client._Client] = None,
     ) -> _Dict:
         """Reference a named Dict, creating if necessary.
 
@@ -666,6 +667,7 @@ class Dict(modal.object.Object):
         namespace=None,
         environment_name: typing.Optional[str] = None,
         create_if_missing: bool = False,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> Dict:
         """Reference a named Dict, creating if necessary.
 

modal/environments.py

@@ -8,6 +8,7 @@ from google.protobuf.wrappers_pb2 import StringValue
 
 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, synchronizer
@@ -52,6 +53,7 @@ class _Environment(_Object, type_prefix="en"):
         name: str,
         *,
         create_if_missing: bool = False,
+        client: Optional[_Client] = None,
     ):
         if name:
             # Allow null names for the case where we want to look up the "default" environment,
@@ -61,7 +63,9 @@ class _Environment(_Object, type_prefix="en"):
             # environments as part of public API when we make this class more useful.
             check_object_name(name, "Environment")
 
-        async def _load(self: _Environment, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _Environment, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             request = api_pb2.EnvironmentGetOrCreateRequest(
                 deployment_name=name,
                 object_creation_type=(
@@ -70,12 +74,17 @@ class _Environment(_Object, type_prefix="en"):
                     else api_pb2.OBJECT_CREATION_TYPE_UNSPECIFIED
                 ),
             )
-            response = await resolver.client.stub.EnvironmentGetOrCreate(request)
+            response = await load_context.client.stub.EnvironmentGetOrCreate(request)
             logger.debug(f"Created environment with id {response.environment_id}")
-            self._hydrate(response.environment_id, resolver.client, response.metadata)
-
-        # TODO environment name (and id?) in the repr? (We should make reprs consistently more useful)
-        return _Environment._from_loader(_load, "Environment()", is_another_app=True, hydrate_lazily=True)
+            self._hydrate(response.environment_id, load_context.client, response.metadata)
+
+        return _Environment._from_loader(
+            _load,
+            f"Environment.from_name({name!r})",
+            is_another_app=True,
+            hydrate_lazily=True,
+            load_context_overrides=LoadContext(client=client),
+        )
 
 
 Environment = synchronize_api(_Environment)
@@ -88,7 +97,7 @@ ENVIRONMENT_CACHE: dict[str, _Environment] = {}
 async def _get_environment_cached(name: str, client: _Client) -> _Environment:
     if name in ENVIRONMENT_CACHE:
         return ENVIRONMENT_CACHE[name]
-    environment = await _Environment.from_name(name).hydrate(client)
+    environment = await _Environment.from_name(name, client=client).hydrate()
     ENVIRONMENT_CACHE[name] = environment
     return environment
 

modal/environments.pyi

@@ -45,7 +45,9 @@ class _Environment(modal._object._Object):
 
     def _hydrate_metadata(self, metadata: google.protobuf.message.Message): ...
     @staticmethod
-    def from_name(name: str, *, create_if_missing: bool = False): ...
+    def from_name(
+        name: str, *, create_if_missing: bool = False, client: typing.Optional[modal.client._Client] = None
+    ): ...
 
 class Environment(modal.object.Object):
     _settings: EnvironmentSettings
@@ -56,7 +58,9 @@ class Environment(modal.object.Object):
 
     def _hydrate_metadata(self, metadata: google.protobuf.message.Message): ...
     @staticmethod
-    def from_name(name: str, *, create_if_missing: bool = False): ...
+    def from_name(
+        name: str, *, create_if_missing: bool = False, client: typing.Optional[modal.client.Client] = None
+    ): ...
 
 async def _get_environment_cached(name: str, client: modal.client._Client) -> _Environment: ...
 

modal/functions.pyi

@@ -1,6 +1,7 @@
 import collections.abc
 import google.protobuf.message
 import modal._functions
+import modal._load_context
 import modal._utils.async_utils
 import modal._utils.function_utils
 import modal.app
@@ -66,7 +67,7 @@ class Function(
     @staticmethod
     def from_local(
         info: modal._utils.function_utils.FunctionInfo,
-        app,
+        app: typing.Optional[modal.app.App],
         image: modal.image.Image,
         env: typing.Optional[dict[str, typing.Optional[str]]] = None,
         secrets: typing.Optional[collections.abc.Collection[modal.secret.Secret]] = None,
@@ -245,10 +246,16 @@ class Function(
     keep_warm: __keep_warm_spec[typing_extensions.Self]
 
     @classmethod
-    def _from_name(cls, app_name: str, name: str, namespace=None, environment_name: typing.Optional[str] = None): ...
+    def _from_name(cls, app_name: str, name: str, *, load_context_overrides: modal._load_context.LoadContext): ...
     @classmethod
     def from_name(
-        cls: type[Function], app_name: str, name: str, *, namespace=None, environment_name: typing.Optional[str] = None
+        cls: type[Function],
+        app_name: str,
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> Function:
         """Reference a Function from a deployed App by its name.
 

modal/image.py

@@ -28,6 +28,7 @@ from typing_extensions import Self
 from modal._serialization import serialize_data_format
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import _Object, live_method_gen
 from ._resolver import Resolver
 from ._serialization import get_preferred_payload_format, serialize
@@ -434,12 +435,16 @@ class _Image(_Object, type_prefix="im"):
 
         base_image = self
 
-        async def _load(self2: "_Image", resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self2: "_Image", resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             self2._hydrate_from_other(base_image)  # same image id as base image as long as it's lazy
             self2._deferred_mounts = tuple(base_image._deferred_mounts) + (mount,)
             self2._serve_mounts = base_image._serve_mounts | ({mount} if mount.is_local() else set())
 
-        img = _Image._from_loader(_load, "Image(local files)", deps=lambda: [base_image, mount])
+        img = _Image._from_loader(
+            _load, "Image(local files)", deps=lambda: [base_image, mount], load_context_overrides=LoadContext.empty()
+        )
         img._added_python_source_set = base_image._added_python_source_set
         return img
 
@@ -523,18 +528,18 @@ class _Image(_Object, type_prefix="im"):
                 deps += (vol,)
             return deps
 
-        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]):
             context_mount = context_mount_function() if context_mount_function else None
             if context_mount:
-                await resolver.load(context_mount)
+                await resolver.load(context_mount, load_context)
 
             if _do_assert_no_mount_layers:
                 for image in base_images.values():
                     # base images can't have
                     image._assert_no_mount_layers()
 
-            assert resolver.app_id  # type narrowing
-            environment = await _get_environment_cached(resolver.environment_name or "", resolver.client)
+            assert load_context.app_id  # type narrowing
+            environment = await _get_environment_cached(load_context.environment_name or "", load_context.client)
             # A bit hacky,but assume that the environment provides a valid builder version
             image_builder_version = cast(ImageBuilderVersion, environment._settings.image_builder_version)
             builder_version = _get_image_builder_version(image_builder_version)
@@ -626,7 +631,7 @@ class _Image(_Object, type_prefix="im"):
             )
 
             req = api_pb2.ImageGetOrCreateRequest(
-                app_id=resolver.app_id,
+                app_id=load_context.app_id,
                 image=image_definition,
                 existing_image_id=existing_object_id or "",  # TODO: ignored
                 build_function_id=build_function_id,
@@ -638,7 +643,7 @@ class _Image(_Object, type_prefix="im"):
                 allow_global_deployment=os.environ.get("MODAL_IMAGE_ALLOW_GLOBAL_DEPLOYMENT") == "1",
                 ignore_cache=config.get("ignore_cache"),
             )
-            resp = await resolver.client.stub.ImageGetOrCreate(req)
+            resp = await load_context.client.stub.ImageGetOrCreate(req)
             image_id = resp.image_id
             result: api_pb2.GenericResult
             metadata: Optional[api_pb2.ImageMetadata] = None
@@ -651,7 +656,7 @@ 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, resolver.client)
+                resp = await _image_await_build_result(image_id, load_context.client)
                 result = resp.result
                 if resp.HasField("metadata"):
                     metadata = resp.metadata
@@ -681,7 +686,7 @@ class _Image(_Object, type_prefix="im"):
             else:
                 raise RemoteError("Unknown status %s!" % result.status)
 
-            self._hydrate(image_id, resolver.client, metadata)
+            self._hydrate(image_id, load_context.client, metadata)
             local_mounts = set()
             for base in base_images.values():
                 local_mounts |= base._serve_mounts
@@ -690,7 +695,7 @@ class _Image(_Object, type_prefix="im"):
             self._serve_mounts = frozenset(local_mounts)
 
         rep = f"Image({dockerfile_function})"
-        obj = _Image._from_loader(_load, rep, deps=_deps)
+        obj = _Image._from_loader(_load, rep, deps=_deps, load_context_overrides=LoadContext.empty())
         obj.force_build = force_build
         obj._added_python_source_set = frozenset.union(
             frozenset(), *(base._added_python_source_set for base in base_images.values())
@@ -863,15 +868,13 @@ class _Image(_Object, type_prefix="im"):
 
         The ID of an Image object can be accessed using `.object_id`.
         """
-        if client is None:
-            client = await _Client.from_env()
 
-        async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
-            resp = await client.stub.ImageFromId(api_pb2.ImageFromIdRequest(image_id=image_id))
-            self._hydrate(resp.image_id, resolver.client, resp.metadata)
+        async def _load(self: _Image, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
+            resp = await load_context.client.stub.ImageFromId(api_pb2.ImageFromIdRequest(image_id=image_id))
+            self._hydrate(resp.image_id, load_context.client, resp.metadata)
 
         rep = f"Image.from_id({image_id!r})"
-        obj = _Image._from_loader(_load, rep)
+        obj = _Image._from_loader(_load, rep, load_context_overrides=LoadContext(client=client))
 
         return obj
 
@@ -930,11 +933,8 @@ class _Image(_Object, type_prefix="im"):
         if app.app_id is None:
             raise InvalidError("App has not been initialized yet. Use the content manager `app.run()` or `App.lookup`")
 
-        app_id = app.app_id
-        app_client = app._client or await _Client.from_env()
-
-        resolver = Resolver(app_client, app_id=app_id)
-        await resolver.load(self)
+        resolver = Resolver()
+        await resolver.load(self, app._root_load_context)
         return self
 
     def pip_install(