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

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

modal/queue.py

@@ -14,6 +14,7 @@ from synchronicity.async_wrap import asynccontextmanager
 
 from modal_proto import api_pb2
 
+from ._load_context import LoadContext
 from ._object import (
     EPHEMERAL_OBJECT_HEARTBEAT_SLEEP,
     _get_environment_name,
@@ -361,6 +362,7 @@ class _Queue(_Object, type_prefix="qu"):
         namespace=None,  # mdmd:line-hidden
         environment_name: Optional[str] = None,
         create_if_missing: bool = False,
+        client: Optional[_Client] = None,
     ) -> "_Queue":
         """Reference a named Queue, creating if necessary.
 
@@ -376,17 +378,24 @@ class _Queue(_Object, type_prefix="qu"):
         check_object_name(name, "Queue")
         warn_if_passing_namespace(namespace, "modal.Queue.from_name")
 
-        async def _load(self: _Queue, resolver: Resolver, existing_object_id: Optional[str]):
+        async def _load(self: _Queue, resolver: Resolver, load_context: LoadContext, existing_object_id: Optional[str]):
             req = api_pb2.QueueGetOrCreateRequest(
                 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),
             )
-            response = await resolver.client.stub.QueueGetOrCreate(req)
-            self._hydrate(response.queue_id, resolver.client, response.metadata)
+            response = await load_context.client.stub.QueueGetOrCreate(req)
+            self._hydrate(response.queue_id, load_context.client, response.metadata)
 
         rep = _Queue._repr(name, environment_name)
-        return _Queue._from_loader(_load, rep, is_another_app=True, hydrate_lazily=True, name=name)
+        return _Queue._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(name: str, *, client: Optional[_Client] = None, environment_name: Optional[str] = None):

modal/queue.pyi

@@ -464,7 +464,12 @@ class _Queue(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,
     ) -> _Queue:
         """Reference a named Queue, creating if necessary.
 
@@ -721,7 +726,12 @@ class Queue(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,
     ) -> Queue:
         """Reference a named Queue, creating if necessary.
 

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
 
@@ -76,6 +76,7 @@ async def _init_local_app_existing(client: _Client, existing_app_id: str, enviro
 async def _init_local_app_new(
     client: _Client,
     description: str,
+    tags: dict[str, str],
     app_state: int,  # ValueType
     environment_name: str = "",
     interactive: bool = False,
@@ -84,6 +85,7 @@ async def _init_local_app_new(
         description=description,
         environment_name=environment_name,
         app_state=app_state,  # type: ignore
+        tags=tags,
     )
     app_resp, _ = await gather_cancel_on_exc(  # TODO: use TaskGroup?
         client.stub.AppCreate(app_req),
@@ -102,6 +104,7 @@ async def _init_local_app_new(
 async def _init_local_app_from_name(
     client: _Client,
     name: str,
+    tags: dict[str, str],
     environment_name: str = "",
 ) -> RunningApp:
     # Look up any existing deployment
@@ -117,23 +120,19 @@ async def _init_local_app_from_name(
         return await _init_local_app_existing(client, existing_app_id, environment_name)
     else:
         return await _init_local_app_new(
-            client, name, api_pb2.APP_STATE_INITIALIZING, environment_name=environment_name
+            client, name, tags, api_pb2.APP_STATE_INITIALIZING, environment_name=environment_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}
@@ -156,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
 
@@ -164,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):
@@ -263,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")
@@ -285,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()
@@ -295,21 +293,25 @@ async def _run_app(
         msg = "Interactive mode requires output to be enabled. (Use the the `modal.enable_output()` context manager.)"
         raise InvalidError(msg)
 
+    local_app_state = app._local_state
+
     running_app: RunningApp = await _init_local_app_new(
-        client,
+        load_context.client,
         app.description or "",
-        environment_name=environment_name or "",
+        local_app_state.tags,
+        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
@@ -330,25 +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
+                )
             )
 
-        local_app_state = app._local_state
         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 = (
@@ -376,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:
@@ -385,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)
@@ -416,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
@@ -444,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,
@@ -493,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 ""
@@ -521,12 +517,25 @@ async def _deploy_app(
     if client is None:
         client = await _Client.from_env()
 
+    local_app_state = app._local_state
     t0 = time.time()
 
     # Get git information to track deployment history
     commit_info_task = asyncio.create_task(get_git_commit_info())
 
-    running_app: RunningApp = await _init_local_app_from_name(client, name, environment_name=environment_name)
+    # 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(
+        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:
         # Start heartbeats loop to keep the client alive
@@ -537,12 +546,7 @@ async def _deploy_app(
 
         try:
             # Create all members
-            await _create_all_objects(
-                client,
-                running_app,
-                app._local_state,
-                environment_name=environment_name,
-            )
+            await _create_all_objects(running_app, local_app_state, root_load_context)
 
             commit_info = None
             try:
@@ -554,7 +558,7 @@ async def _deploy_app(
                 client,
                 running_app,
                 api_pb2.APP_STATE_DEPLOYED,
-                app._local_state,
+                local_app_state,
                 name=name,
                 deployment_tag=tag,
                 commit_info=commit_info,

modal/runner.pyi

@@ -1,3 +1,4 @@
+import modal._load_context
 import modal.app
 import modal.client
 import modal.running_app
@@ -16,18 +17,18 @@ async def _init_local_app_existing(
 async def _init_local_app_new(
     client: modal.client._Client,
     description: str,
+    tags: dict[str, str],
     app_state: int,
     environment_name: str = "",
     interactive: bool = False,
 ) -> modal.running_app.RunningApp: ...
 async def _init_local_app_from_name(
-    client: modal.client._Client, name: str, environment_name: str = ""
+    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