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

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(

modal/mount.py

@@ -20,7 +20,8 @@ import modal.file_pattern_matcher
 from modal_proto import api_pb2
 from modal_version import __version__
 
-from ._object import _get_environment_name, _Object
+from ._load_context import LoadContext
+from ._object import _Object
 from ._resolver import Resolver
 from ._utils.async_utils import TaskContext, aclosing, async_map, synchronize_api
 from ._utils.blob_utils import FileUploadSpec, blob_upload_file, get_file_upload_spec_from_path
@@ -310,7 +311,7 @@ class _Mount(_Object, type_prefix="mo"):
     _entries: Optional[list[_MountEntry]] = None
     _deployment_name: Optional[str] = None
     _namespace: Optional[int] = None
-    _environment_name: Optional[str] = None
+
     _allow_overwrite: bool = False
     _content_checksum_sha256_hex: Optional[str] = None
 
@@ -325,7 +326,12 @@ class _Mount(_Object, type_prefix="mo"):
                 return None
             return (_Mount._type_prefix, "local", frozenset(included_files))
 
-        obj = _Mount._from_loader(_Mount._load_mount, rep, deduplication_key=mount_content_deduplication_key)
+        obj = _Mount._from_loader(
+            _Mount._load_mount,
+            rep,
+            deduplication_key=mount_content_deduplication_key,
+            load_context_overrides=LoadContext.empty(),
+        )
         obj._entries = entries
         obj._is_local = True
         return obj
@@ -477,6 +483,7 @@ class _Mount(_Object, type_prefix="mo"):
     async def _load_mount(
         self: "_Mount",
         resolver: Resolver,
+        load_context: LoadContext,
         existing_object_id: Optional[str],
     ):
         t0 = time.monotonic()
@@ -518,7 +525,7 @@ class _Mount(_Object, type_prefix="mo"):
 
             request = api_pb2.MountPutFileRequest(sha256_hex=file_spec.sha256_hex)
             accounted_hashes.add(file_spec.sha256_hex)
-            response = await resolver.client.stub.MountPutFile(request, retry=Retry(base_delay=1))
+            response = await load_context.client.stub.MountPutFile(request, retry=Retry(base_delay=1))
 
             if response.exists:
                 n_finished += 1
@@ -532,7 +539,7 @@ class _Mount(_Object, type_prefix="mo"):
                 async with blob_upload_concurrency:
                     with file_spec.source() as fp:
                         blob_id = await blob_upload_file(
-                            fp, resolver.client.stub, sha256_hex=file_spec.sha256_hex, md5_hex=file_spec.md5_hex
+                            fp, load_context.client.stub, sha256_hex=file_spec.sha256_hex, md5_hex=file_spec.md5_hex
                         )
                 logger.debug(f"Uploading blob file {file_spec.source_description} as {remote_filename}")
                 request2 = api_pb2.MountPutFileRequest(data_blob_id=blob_id, sha256_hex=file_spec.sha256_hex)
@@ -544,7 +551,7 @@ class _Mount(_Object, type_prefix="mo"):
 
             start_time = time.monotonic()
             while time.monotonic() - start_time < MOUNT_PUT_FILE_CLIENT_TIMEOUT:
-                response = await resolver.client.stub.MountPutFile(request2, retry=Retry(base_delay=1))
+                response = await load_context.client.stub.MountPutFile(request2, retry=Retry(base_delay=1))
                 if response.exists:
                     n_finished += 1
                     return mount_file
@@ -552,7 +559,7 @@ class _Mount(_Object, type_prefix="mo"):
             raise modal.exception.MountUploadTimeoutError(f"Mounting of {file_spec.source_description} timed out")
 
         # Upload files, or check if they already exist.
-        n_concurrent_uploads = 512
+        n_concurrent_uploads = 64
         files: list[api_pb2.MountFile] = []
         async with aclosing(
             async_map(_Mount._get_files(self._entries), _put_file, concurrency=n_concurrent_uploads)
@@ -574,28 +581,28 @@ class _Mount(_Object, type_prefix="mo"):
             req = api_pb2.MountGetOrCreateRequest(
                 deployment_name=self._deployment_name,
                 namespace=self._namespace,
-                environment_name=self._environment_name,
+                environment_name=load_context.environment_name,
                 object_creation_type=creation_type,
                 files=files,
             )
-        elif resolver.app_id is not None:
+        elif load_context.app_id is not None:
             req = api_pb2.MountGetOrCreateRequest(
                 object_creation_type=api_pb2.OBJECT_CREATION_TYPE_ANONYMOUS_OWNED_BY_APP,
                 files=files,
-                app_id=resolver.app_id,
+                app_id=load_context.app_id,
             )
         else:
             req = api_pb2.MountGetOrCreateRequest(
                 object_creation_type=api_pb2.OBJECT_CREATION_TYPE_EPHEMERAL,
                 files=files,
-                environment_name=resolver.environment_name,
+                environment_name=load_context.environment_name,
             )
 
-        resp = await resolver.client.stub.MountGetOrCreate(req, retry=Retry(base_delay=1))
+        resp = await load_context.client.stub.MountGetOrCreate(req, retry=Retry(base_delay=1))
         status_row.finish(f"Created mount {message_label}")
 
         logger.debug(f"Uploaded {total_uploads} new files and {total_bytes} bytes in {time.monotonic() - t0}s")
-        self._hydrate(resp.mount_id, resolver.client, resp.handle_metadata)
+        self._hydrate(resp.mount_id, load_context.client, resp.handle_metadata)
 
     @staticmethod
     def _from_local_python_packages(
@@ -628,19 +635,25 @@ class _Mount(_Object, type_prefix="mo"):
         *,
         namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
         environment_name: Optional[str] = None,
+        client: Optional[_Client] = None,
     ) -> "_Mount":
         """mdmd:hidden"""
 
-        async def _load(provider: _Mount, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(provider: _Mount, resolver: Resolver, load_context, existing_object_id: Optional[str]):
             req = api_pb2.MountGetOrCreateRequest(
                 deployment_name=name,
                 namespace=namespace,
-                environment_name=_get_environment_name(environment_name, resolver),
+                environment_name=load_context.environment_name,
             )
-            response = await resolver.client.stub.MountGetOrCreate(req)
-            provider._hydrate(response.mount_id, resolver.client, response.handle_metadata)
-
-        return _Mount._from_loader(_load, "Mount()", hydrate_lazily=True)
+            response = await load_context.client.stub.MountGetOrCreate(req)
+            provider._hydrate(response.mount_id, load_context.client, response.handle_metadata)
+
+        return _Mount._from_loader(
+            _load,
+            "Mount()",
+            hydrate_lazily=True,
+            load_context_overrides=LoadContext(environment_name=environment_name, client=client),
+        )
 
     async def _deploy(
         self: "_Mount",
@@ -652,15 +665,12 @@ class _Mount(_Object, type_prefix="mo"):
         client: Optional[_Client] = None,
     ) -> None:
         check_object_name(deployment_name, "Mount")
-        environment_name = _get_environment_name(environment_name, resolver=None)
         self._deployment_name = deployment_name
         self._namespace = namespace
-        self._environment_name = environment_name
         self._allow_overwrite = allow_overwrite
-        if client is None:
-            client = await _Client.from_env()
-        resolver = Resolver(client=client, environment_name=environment_name)
-        await resolver.load(self)
+        resolver = Resolver()
+        root_metadata = LoadContext(client=client, environment_name=environment_name)
+        await resolver.load(self, root_metadata)
 
     def _get_metadata(self) -> api_pb2.MountHandleMetadata:
         if self._content_checksum_sha256_hex is None:

modal/mount.pyi

@@ -1,5 +1,6 @@
 import collections.abc
 import google.protobuf.message
+import modal._load_context
 import modal._object
 import modal._resolver
 import modal._utils.blob_utils
@@ -154,7 +155,6 @@ class _Mount(modal._object._Object):
     _entries: typing.Optional[list[_MountEntry]]
     _deployment_name: typing.Optional[str]
     _namespace: typing.Optional[int]
-    _environment_name: typing.Optional[str]
     _allow_overwrite: bool
     _content_checksum_sha256_hex: typing.Optional[str]
 
@@ -216,7 +216,10 @@ class _Mount(modal._object._Object):
         entries: list[_MountEntry],
     ) -> collections.abc.AsyncGenerator[modal._utils.blob_utils.FileUploadSpec, None]: ...
     async def _load_mount(
-        self: _Mount, resolver: modal._resolver.Resolver, existing_object_id: typing.Optional[str]
+        self: _Mount,
+        resolver: modal._resolver.Resolver,
+        load_context: modal._load_context.LoadContext,
+        existing_object_id: typing.Optional[str],
     ): ...
     @staticmethod
     def _from_local_python_packages(
@@ -226,7 +229,13 @@ class _Mount(modal._object._Object):
         ignore: typing.Union[typing.Sequence[str], collections.abc.Callable[[pathlib.Path], bool], None] = None,
     ) -> _Mount: ...
     @staticmethod
-    def from_name(name: str, *, namespace=1, environment_name: typing.Optional[str] = None) -> _Mount:
+    def from_name(
+        name: str,
+        *,
+        namespace=1,
+        environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client._Client] = None,
+    ) -> _Mount:
         """mdmd:hidden"""
         ...
 
@@ -269,7 +278,6 @@ class Mount(modal.object.Object):
     _entries: typing.Optional[list[_MountEntry]]
     _deployment_name: typing.Optional[str]
     _namespace: typing.Optional[int]
-    _environment_name: typing.Optional[str]
     _allow_overwrite: bool
     _content_checksum_sha256_hex: typing.Optional[str]
 
@@ -342,8 +350,20 @@ class Mount(modal.object.Object):
     _get_files: ___get_files_spec
 
     class ___load_mount_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /, resolver: modal._resolver.Resolver, existing_object_id: typing.Optional[str]): ...
-        async def aio(self, /, resolver: modal._resolver.Resolver, existing_object_id: typing.Optional[str]): ...
+        def __call__(
+            self,
+            /,
+            resolver: modal._resolver.Resolver,
+            load_context: modal._load_context.LoadContext,
+            existing_object_id: typing.Optional[str],
+        ): ...
+        async def aio(
+            self,
+            /,
+            resolver: modal._resolver.Resolver,
+            load_context: modal._load_context.LoadContext,
+            existing_object_id: typing.Optional[str],
+        ): ...
 
     _load_mount: ___load_mount_spec[typing_extensions.Self]
 
@@ -355,7 +375,13 @@ class Mount(modal.object.Object):
         ignore: typing.Union[typing.Sequence[str], collections.abc.Callable[[pathlib.Path], bool], None] = None,
     ) -> Mount: ...
     @staticmethod
-    def from_name(name: str, *, namespace=1, environment_name: typing.Optional[str] = None) -> Mount:
+    def from_name(
+        name: str,
+        *,
+        namespace=1,
+        environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client.Client] = None,
+    ) -> Mount:
         """mdmd:hidden"""
         ...
 

modal/network_file_system.py

@@ -11,6 +11,7 @@ from synchronicity.async_wrap import asynccontextmanager
 import modal
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import (
     EPHEMERAL_OBJECT_HEARTBEAT_SLEEP,
     _get_environment_name,
@@ -95,6 +96,7 @@ class _NetworkFileSystem(_Object, type_prefix="sv"):
         namespace=None,  # mdmd:line-hidden
         environment_name: Optional[str] = None,
         create_if_missing: bool = False,
+        client: Optional[_Client] = None,
     ) -> "_NetworkFileSystem":
         """Reference a NetworkFileSystem by its name, creating if necessary.
 
@@ -113,15 +115,17 @@ class _NetworkFileSystem(_Object, type_prefix="sv"):
         check_object_name(name, "NetworkFileSystem")
         warn_if_passing_namespace(namespace, "modal.NetworkFileSystem.from_name")
 
-        async def _load(self: _NetworkFileSystem, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(
+            self: _NetworkFileSystem, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]
+        ):
             req = api_pb2.SharedVolumeGetOrCreateRequest(
                 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),
             )
             try:
-                response = await resolver.client.stub.SharedVolumeGetOrCreate(req)
-                self._hydrate(response.shared_volume_id, resolver.client, None)
+                response = await load_context.client.stub.SharedVolumeGetOrCreate(req)
+                self._hydrate(response.shared_volume_id, load_context.client, None)
             except modal.exception.NotFoundError as exc:
                 if exc.args[0] == "App has wrong entity vo":
                     raise InvalidError(
@@ -129,7 +133,12 @@ class _NetworkFileSystem(_Object, type_prefix="sv"):
                     )
                 raise
 
-        return _NetworkFileSystem._from_loader(_load, "NetworkFileSystem()", hydrate_lazily=True)
+        return _NetworkFileSystem._from_loader(
+            _load,
+            "NetworkFileSystem()",
+            hydrate_lazily=True,
+            load_context_overrides=LoadContext(environment_name=environment_name, client=client),
+        )
 
     @classmethod
     @asynccontextmanager

modal/network_file_system.pyi

@@ -54,7 +54,12 @@ class _NetworkFileSystem(modal._object._Object):
     """
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, create_if_missing: bool = False
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        create_if_missing: bool = False,
+        client: typing.Optional[modal.client._Client] = None,
     ) -> _NetworkFileSystem:
         """Reference a NetworkFileSystem by its name, creating if necessary.
 
@@ -210,7 +215,12 @@ class NetworkFileSystem(modal.object.Object):
 
     @staticmethod
     def from_name(
-        name: str, *, namespace=None, environment_name: typing.Optional[str] = None, create_if_missing: bool = False
+        name: str,
+        *,
+        namespace=None,
+        environment_name: typing.Optional[str] = None,
+        create_if_missing: bool = False,
+        client: typing.Optional[modal.client.Client] = None,
     ) -> NetworkFileSystem:
         """Reference a NetworkFileSystem by its name, creating if necessary.
 

modal/object.pyi

@@ -1,5 +1,6 @@
 import collections.abc
 import google.protobuf.message
+import modal._load_context
 import modal._resolver
 import modal.client
 import typing
@@ -12,12 +13,14 @@ class Object:
     _prefix_to_type: typing.ClassVar[dict[str, type]]
     _load: typing.Optional[
         collections.abc.Callable[
-            [typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+            [typing_extensions.Self, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+            collections.abc.Awaitable[None],
         ]
     ]
     _preload: typing.Optional[
         collections.abc.Callable[
-            [typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+            [typing_extensions.Self, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+            collections.abc.Awaitable[None],
         ]
     ]
     _rep: str
@@ -27,6 +30,7 @@ class Object:
     _deduplication_key: typing.Optional[
         collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Hashable]]
     ]
+    _load_context_overrides: modal._load_context.LoadContext
     _object_id: typing.Optional[str]
     _client: typing.Optional[modal.client.Client]
     _is_hydrated: bool
@@ -46,16 +50,22 @@ class Object:
             /,
             rep: str,
             load: typing.Optional[
-                collections.abc.Callable[[SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], None]
+                collections.abc.Callable[
+                    [SUPERSELF, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]], None
+                ]
             ] = None,
             is_another_app: bool = False,
             preload: typing.Optional[
-                collections.abc.Callable[[SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], None]
+                collections.abc.Callable[
+                    [SUPERSELF, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]], None
+                ]
             ] = None,
             hydrate_lazily: bool = False,
             deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
             deduplication_key: typing.Optional[collections.abc.Callable[[], collections.abc.Hashable]] = None,
             name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: typing.Optional[modal._load_context.LoadContext] = None,
         ): ...
         def aio(
             self,
@@ -63,13 +73,15 @@ class Object:
             rep: str,
             load: typing.Optional[
                 collections.abc.Callable[
-                    [SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+                    [SUPERSELF, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+                    collections.abc.Awaitable[None],
                 ]
             ] = None,
             is_another_app: bool = False,
             preload: typing.Optional[
                 collections.abc.Callable[
-                    [SUPERSELF, modal._resolver.Resolver, typing.Optional[str]], collections.abc.Awaitable[None]
+                    [SUPERSELF, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+                    collections.abc.Awaitable[None],
                 ]
             ] = None,
             hydrate_lazily: bool = False,
@@ -78,6 +90,8 @@ class Object:
                 collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Hashable]]
             ] = None,
             name: typing.Optional[str] = None,
+            *,
+            load_context_overrides: typing.Optional[modal._load_context.LoadContext] = None,
         ): ...
 
     _init: ___init_spec[typing_extensions.Self]
@@ -101,16 +115,29 @@ class Object:
     @classmethod
     def _from_loader(
         cls,
-        load: collections.abc.Callable[[typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], None],
+        load: collections.abc.Callable[
+            [typing_extensions.Self, modal._resolver.Resolver, modal._load_context.LoadContext, typing.Optional[str]],
+            None,
+        ],
         rep: str,
         is_another_app: bool = False,
         preload: typing.Optional[
-            collections.abc.Callable[[typing_extensions.Self, modal._resolver.Resolver, typing.Optional[str]], None]
+            collections.abc.Callable[
+                [
+                    typing_extensions.Self,
+                    modal._resolver.Resolver,
+                    modal._load_context.LoadContext,
+                    typing.Optional[str],
+                ],
+                None,
+            ]
         ] = None,
         hydrate_lazily: bool = False,
         deps: typing.Optional[collections.abc.Callable[..., collections.abc.Sequence[Object]]] = None,
         deduplication_key: typing.Optional[collections.abc.Callable[[], collections.abc.Hashable]] = None,
         name: typing.Optional[str] = None,
+        *,
+        load_context_overrides: modal._load_context.LoadContext,
     ): ...
     @staticmethod
     def _get_type_from_id(object_id: str) -> type[Object]: ...

modal/proxy.py

@@ -3,9 +3,11 @@ from typing import Optional
 
 from modal_proto import api_pb2
 
-from ._object import _get_environment_name, _Object
+from ._load_context import LoadContext
+from ._object import _Object
 from ._resolver import Resolver
 from ._utils.async_utils import synchronize_api
+from .client import _Client
 
 
 class _Proxy(_Object, type_prefix="pr"):
@@ -20,6 +22,7 @@ class _Proxy(_Object, type_prefix="pr"):
         name: str,
         *,
         environment_name: Optional[str] = None,
+        client: Optional[_Client] = None,
     ) -> "_Proxy":
         """Reference a Proxy by its name.
 
@@ -28,16 +31,21 @@ class _Proxy(_Object, type_prefix="pr"):
 
         """
 
-        async def _load(self: _Proxy, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: _Proxy, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             req = api_pb2.ProxyGetRequest(
                 name=name,
-                environment_name=_get_environment_name(environment_name, resolver),
+                environment_name=load_context.environment_name,
             )
-            response: api_pb2.ProxyGetResponse = await resolver.client.stub.ProxyGet(req)
-            self._hydrate(response.proxy.proxy_id, resolver.client, None)
+            response: api_pb2.ProxyGetResponse = await load_context.client.stub.ProxyGet(req)
+            self._hydrate(response.proxy.proxy_id, load_context.client, None)
 
         rep = _Proxy._repr(name, environment_name)
-        return _Proxy._from_loader(_load, rep, is_another_app=True)
+        return _Proxy._from_loader(
+            _load,
+            rep,
+            is_another_app=True,
+            load_context_overrides=LoadContext(client=client, environment_name=environment_name),
+        )
 
 
 Proxy = synchronize_api(_Proxy, target_module=__name__)

modal/proxy.pyi

@@ -1,4 +1,5 @@
 import modal._object
+import modal.client
 import modal.object
 import typing
 
@@ -9,7 +10,12 @@ class _Proxy(modal._object._Object):
     a database. See [the guide](https://modal.com/docs/guide/proxy-ips) for more information.
     """
     @staticmethod
-    def from_name(name: str, *, environment_name: typing.Optional[str] = None) -> _Proxy:
+    def from_name(
+        name: str,
+        *,
+        environment_name: typing.Optional[str] = None,
+        client: typing.Optional[modal.client._Client] = None,
+    ) -> _Proxy:
         """Reference a Proxy by its name.
 
         In contrast to most other Modal objects, new Proxy objects must be
@@ -28,7 +34,9 @@ class Proxy(modal.object.Object):
         ...
 
     @staticmethod
-    def from_name(name: str, *, environment_name: typing.Optional[str] = None) -> Proxy:
+    def from_name(
+        name: str, *, environment_name: typing.Optional[str] = None, client: typing.Optional[modal.client.Client] = None
+    ) -> Proxy:
         """Reference a Proxy by its name.
 
         In contrast to most other Modal objects, new Proxy objects must be