modal 1.0.6.dev58__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/runner.py

@@ -8,8 +8,8 @@ import asyncio
 import dataclasses
 import os
 import time
-import typing
 from collections.abc import AsyncGenerator
+from contextlib import nullcontext
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, Optional, TypeVar
 
@@ -18,6 +18,8 @@ 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
 
 from ._functions import _Function
@@ -28,13 +30,12 @@ from ._traceback import print_server_warnings, traceback_contains_remote_call
 from ._utils.async_utils import TaskContext, gather_cancel_on_exc, synchronize_api
 from ._utils.deprecation import warn_if_passing_namespace
 from ._utils.git_utils import get_git_commit_info
-from ._utils.grpc_utils import retry_transient_errors
 from ._utils.name_utils import check_object_name, is_valid_tag
 from .client import HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT, _Client
 from .cls import _Cls
 from .config import config, logger
 from .environments import _get_environment_cached
-from .exception import InteractiveTimeoutError, InvalidError, RemoteError, _CliUserExecutionError
+from .exception import ConnectionError, InteractiveTimeoutError, InvalidError, RemoteError, _CliUserExecutionError
 from .output import _get_output_manager, enable_output
 from .running_app import RunningApp, running_app_from_layout
 from .sandbox import _Sandbox
@@ -42,9 +43,7 @@ from .secret import _Secret
 from .stream_type import StreamType
 
 if TYPE_CHECKING:
-    from .app import _App
-else:
-    _App = TypeVar("_App")
+    import modal.app
 
 
 V = TypeVar("V")
@@ -55,14 +54,14 @@ async def _heartbeat(client: _Client, app_id: str) -> None:
     # TODO(erikbern): we should capture exceptions here
     # * if request fails: destroy the client
     # * if server says the app is gone: print a helpful warning about detaching
-    await retry_transient_errors(client.stub.AppHeartbeat, request, attempt_timeout=HEARTBEAT_TIMEOUT)
+    await client.stub.AppHeartbeat(request, retry=Retry(attempt_timeout=HEARTBEAT_TIMEOUT))
 
 
 async def _init_local_app_existing(client: _Client, existing_app_id: str, environment_name: str) -> RunningApp:
     # Get all the objects first
     obj_req = api_pb2.AppGetLayoutRequest(app_id=existing_app_id)
     obj_resp, _ = await gather_cancel_on_exc(
-        retry_transient_errors(client.stub.AppGetLayout, obj_req),
+        client.stub.AppGetLayout(obj_req),
         # Cache the environment associated with the app now as we will use it later
         _get_environment_cached(environment_name, client),
     )
@@ -77,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,
@@ -85,9 +85,10 @@ 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?
-        retry_transient_errors(client.stub.AppCreate, app_req),
+        client.stub.AppCreate(app_req),
         # Cache the environment associated with the app now as we will use it later
         _get_environment_cached(environment_name, client),
     )
@@ -103,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
@@ -110,7 +112,7 @@ async def _init_local_app_from_name(
         name=name,
         environment_name=environment_name,
     )
-    app_resp = await retry_transient_errors(client.stub.AppGetByDeploymentName, app_req)
+    app_resp = await client.stub.AppGetByDeploymentName(app_req)
     existing_app_id = app_resp.app_id or None
 
     # Grab the app
@@ -118,24 +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,
-    functions: dict[str, _Function],
-    classes: dict[str, _Cls],
-    environment_name: str,
+    local_app_state: "modal.app._LocalAppState",
+    load_context: LoadContext,
 ) -> None:
     """Create objects that have been defined but not created on the server."""
-    indexed_objects: dict[str, _Object] = {**functions, **classes}
-    resolver = Resolver(
-        client,
-        environment_name=environment_name,
-        app_id=running_app.app_id,
-    )
+    indexed_objects: dict[str, _Object] = {**local_app_state.functions, **local_app_state.classes}
+
+    resolver = Resolver()
     with resolver.display():
         # Get current objects, and reset all objects
         tag_to_object_id = {**running_app.function_ids, **running_app.class_ids}
@@ -158,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
 
@@ -166,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):
@@ -181,29 +179,29 @@ async def _publish_app(
     client: _Client,
     running_app: RunningApp,
     app_state: int,  # api_pb2.AppState.value
-    functions: dict[str, _Function],
-    classes: dict[str, _Cls],
-    name: str = "",  # Only relevant for deployments
-    tag: str = "",  # Only relevant for deployments
+    app_local_state: "modal.app._LocalAppState",
+    name: str = "",
+    deployment_tag: str = "",  # Only relevant for deployments
     commit_info: Optional[api_pb2.CommitInfo] = None,  # Git commit information
 ) -> tuple[str, list[api_pb2.Warning]]:
     """Wrapper for AppPublish RPC."""
-
+    functions = app_local_state.functions
     definition_ids = {obj.object_id: obj._get_metadata().definition_id for obj in functions.values()}  # type: ignore
 
     request = api_pb2.AppPublishRequest(
         app_id=running_app.app_id,
         name=name,
-        deployment_tag=tag,
+        tags=app_local_state.tags,
+        deployment_tag=deployment_tag,
+        commit_info=commit_info,
         app_state=app_state,  # type: ignore  : should be a api_pb2.AppState.value
         function_ids=running_app.function_ids,
         class_ids=running_app.class_ids,
         definition_ids=definition_ids,
-        commit_info=commit_info,
     )
 
     try:
-        response = await retry_transient_errors(client.stub.AppPublish, request)
+        response = await client.stub.AppPublish(request)
     except GRPCError as exc:
         if exc.status == Status.INVALID_ARGUMENT or exc.status == Status.FAILED_PRECONDITION:
             raise InvalidError(exc.message)
@@ -227,7 +225,7 @@ async def _disconnect(
 
     logger.debug("Sending app disconnect/stop request")
     req_disconnect = api_pb2.AppClientDisconnectRequest(app_id=app_id, reason=reason, exception=exc_str)
-    await retry_transient_errors(client.stub.AppClientDisconnect, req_disconnect)
+    await client.stub.AppClientDisconnect(req_disconnect)
     logger.debug("App disconnected")
 
 
@@ -257,16 +255,17 @@ async def _status_based_disconnect(client: _Client, app_id: str, exc_info: Optio
 
 @asynccontextmanager
 async def _run_app(
-    app: _App,
+    app: "modal.app._App",
     *,
     client: Optional[_Client] = None,
     detach: bool = False,
     environment_name: Optional[str] = None,
     interactive: bool = False,
-) -> AsyncGenerator[_App, None]:
+) -> 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")
@@ -287,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()
@@ -297,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
@@ -332,26 +332,37 @@ 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, app._functions, app._classes, environment_name)
+            await _create_all_objects(running_app, local_app_state, load_context)
 
             # Publish the app
-            await _publish_app(client, running_app, app_state, app._functions, app._classes)
+            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 = (
+            "The detached App will keep running. You can track its progress on the Dashboard: "
+            f"[magenta underline]{running_app.app_page_url}[/magenta underline]"
+            "\n\nStream App logs:\n"
+            f"[green]modal app logs {running_app.app_id}[/green]"
+            "\n\nStop the App:\n"
+            f"[green]modal app stop {running_app.app_id}[/green]"
+        )
+
         try:
             # Show logs from dynamically created images.
             # TODO: better way to do this
@@ -360,32 +371,30 @@ async def _run_app(
 
             # Yield to context
             if output_mgr := _get_output_manager():
-                with output_mgr.show_status_spinner():
+                # Don't show status spinner in interactive mode to avoid interfering with breakpoints
+                spinner_ctx = nullcontext() if interactive else output_mgr.show_status_spinner()
+                with spinner_ctx:
                     yield app
             else:
                 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:
                 if output_mgr := _get_output_manager():
                     output_mgr.print(output_mgr.step_completed("Shutting down Modal client."))
-                    output_mgr.print(
-                        "The detached app keeps running. You can track its progress at: "
-                        f"[magenta]{running_app.app_page_url}[/magenta]"
-                        ""
-                    )
+                    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)
@@ -400,9 +409,17 @@ async def _run_app(
                         )
                     )
             return
+        except ConnectionError as e:
+            # If we lose connection to the server after a detached App has started running, it will continue
+            # I think we can only exit "nicely" if we are able to print output though, otherwise we should raise
+            if detach and (output_mgr := _get_output_manager()):
+                output_mgr.print(":white_exclamation_mark: Connection lost!")
+                output_mgr.print(detached_disconnect_msg)
+                return
+            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
@@ -423,28 +440,28 @@ async def _run_app(
 
 
 async def _serve_update(
-    app: _App,
+    app: "modal.app._App",
     existing_app_id: str,
     is_ready: Event,
     environment_name: str,
 ) -> 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,
-            app._functions,
-            app._classes,
-            environment_name,
-        )
+        await _create_all_objects(running_app, local_app_state, load_context)
 
         # Publish the updated app
-        await _publish_app(client, running_app, api_pb2.APP_STATE_UNSPECIFIED, app._functions, app._classes)
+        await _publish_app(
+            load_context.client,
+            running_app,
+            app_state=api_pb2.APP_STATE_UNSPECIFIED,
+            app_local_state=local_app_state,
+        )
 
         # Communicate to the parent process
         is_ready.set()
@@ -464,7 +481,7 @@ class DeployResult:
 
 
 async def _deploy_app(
-    app: _App,
+    app: "modal.app._App",
     name: Optional[str] = None,
     namespace: Any = None,  # mdmd:line-hidden
     client: Optional[_Client] = None,
@@ -475,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 ""
@@ -493,7 +507,7 @@ async def _deploy_app(
     else:
         check_object_name(name, "App")
 
-    if tag and not is_valid_tag(tag):
+    if tag and not is_valid_tag(tag, max_length=50):
         raise InvalidError(
             f"Deployment tag {tag!r} is invalid."
             "\n\nTags may only contain alphanumeric characters, dashes, periods, and underscores, "
@@ -503,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
@@ -519,13 +546,7 @@ async def _deploy_app(
 
         try:
             # Create all members
-            await _create_all_objects(
-                client,
-                running_app,
-                app._functions,
-                app._classes,
-                environment_name=environment_name,
-            )
+            await _create_all_objects(running_app, local_app_state, root_load_context)
 
             commit_info = None
             try:
@@ -537,11 +558,10 @@ async def _deploy_app(
                 client,
                 running_app,
                 api_pb2.APP_STATE_DEPLOYED,
-                app._functions,
-                app._classes,
-                name,
-                tag,
-                commit_info,
+                local_app_state,
+                name=name,
+                deployment_tag=tag,
+                commit_info=commit_info,
             )
         except Exception as e:
             # Note that AppClientDisconnect only stops the app if it's still initializing, and is a no-op otherwise.
@@ -561,7 +581,7 @@ async def _deploy_app(
 
 
 async def _interactive_shell(
-    _app: _App, cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: Any
+    _app: "modal.app._App", cmds: list[str], environment_name: str = "", pty: bool = True, **kwargs: Any
 ) -> None:
     """Run an interactive shell (like `bash`) within the image for this app.
 
@@ -607,19 +627,19 @@ async def _interactive_shell(
 
         try:
             if pty:
-                container_process = await sandbox.exec(
+                container_process = await sandbox._exec(
                     *sandbox_cmds, pty_info=get_pty_info(shell=True) if pty else None
                 )
                 await container_process.attach()
             else:
-                container_process = await sandbox.exec(
+                container_process = await sandbox._exec(
                     *sandbox_cmds, stdout=StreamType.STDOUT, stderr=StreamType.STDOUT
                 )
                 await container_process.wait()
         except InteractiveTimeoutError:
             # Check on status of Sandbox. It may have crashed, causing connection failure.
             req = api_pb2.SandboxWaitRequest(sandbox_id=sandbox._object_id, timeout=0)
-            resp = await retry_transient_errors(sandbox._client.stub.SandboxWait, req)
+            resp = await sandbox._client.stub.SandboxWait(req)
             if resp.result.exception:
                 raise RemoteError(resp.result.exception)
             else: