modAL 1.5.1__tar.gz → 1.5.1.dev0__tar.gz

{modal-1.5.1 → modal-1.5.1.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.5.1
+Version: 1.5.1.dev0
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License-Expression: Apache-2.0
@@ -23,15 +23,12 @@ Requires-Dist: grpclib<0.4.10,>=0.4.7; python_version < "3.14"
 Requires-Dist: grpclib<0.4.10,>=0.4.9; python_version >= "3.14"
 Requires-Dist: protobuf!=4.24.0,<7.0,>=3.19
 Requires-Dist: rich>=12.0.0
-Requires-Dist: synchronicity~=0.12.5
+Requires-Dist: synchronicity~=0.12.3
 Requires-Dist: toml
 Requires-Dist: types-certifi
 Requires-Dist: types-toml
 Requires-Dist: watchfiles
 Requires-Dist: typing_extensions~=4.6
-Provides-Extra: api-proxy-support
-Requires-Dist: python-socks[asyncio]>=2.5; extra == "api-proxy-support"
-Requires-Dist: aiohttp-socks>=0.8; extra == "api-proxy-support"
 Dynamic: license-file
 
 # Modal Python SDK
@@ -48,7 +45,7 @@ access to serverless cloud compute from Python scripts on your local computer.
 See the [online documentation](https://modal.com/docs/guide) for many
 [example applications](https://modal.com/docs/examples),
 a [user guide](https://modal.com/docs/guide), and the detailed
-[API reference](https://modal.com/docs/sdk/py/latest).
+[API reference](https://modal.com/docs/reference).
 
 ## Installation
 

{modal-1.5.1 → modal-1.5.1.dev0}/README.md

@@ -12,7 +12,7 @@ access to serverless cloud compute from Python scripts on your local computer.
 See the [online documentation](https://modal.com/docs/guide) for many
 [example applications](https://modal.com/docs/examples),
 a [user guide](https://modal.com/docs/guide), and the detailed
-[API reference](https://modal.com/docs/sdk/py/latest).
+[API reference](https://modal.com/docs/reference).
 
 ## Installation
 

{modal-1.5.1 → modal-1.5.1.dev0}/modal/__init__.py

@@ -43,7 +43,6 @@ try:
     from .schedule import Cron, Period
     from .scheduler_placement import SchedulerPlacement
     from .secret import Secret
-    from .server import Server
     from .snapshot import SandboxSnapshot
     from .volume import Volume
     from .workspace import Workspace
@@ -81,7 +80,6 @@ __all__ = [
     "SandboxSnapshot",
     "SchedulerPlacement",
     "Secret",
-    "Server",
     "Tunnel",
     "Volume",
     "Workspace",

modal-1.5.1.dev0/modal/_billing.py

@@ -0,0 +1,94 @@
+# Copyright Modal Labs 2025
+from datetime import datetime, timezone
+from decimal import Decimal
+from typing import TypedDict
+
+from modal_proto import api_pb2
+
+from .client import _Client
+from .exception import InvalidError
+
+
+class WorkspaceBillingReportItem(TypedDict):
+    object_id: str
+    description: str
+    environment_name: str
+    interval_start: datetime
+    cost: Decimal
+    tags: dict[str, str]
+
+
+async def _workspace_billing_report(
+    *,
+    start: datetime,  # Start of the report, inclusive
+    end: datetime | None = None,  # End of the report, exclusive
+    resolution: str = "d",  # Resolution, e.g. "d" for daily or "h" for hourly
+    tag_names: list[str] | None = None,  # Optional additional metadata to include
+    client: _Client | None = None,
+) -> list[WorkspaceBillingReportItem]:
+    """Generate a tabular report of workspace usage by object and time.
+
+    The result will be a list of dictionaries for each interval (determined by `resolution`)
+    between the `start` and `end` limits. The dictionary represents a single Modal object
+    that billing can be attributed to (e.g., an App) along with metadata (including user-defined
+    tags) for identifying that object.
+
+    The `start` and `end` parameters are required to either have a UTC timezone or to be
+    timezone-naive (which will be interpreted as UTC times). The timestamps in the result will
+    be in UTC. Cost will be reported for full intervals, even if the provided `start` or `end`
+    parameters are partial: `start` will be rounded to the beginning of its interval, while
+    partial `end` intervals will be excluded.
+
+    Additional user-provided metadata can be included in the report if the objects have tags
+    and `tag_names` (i.e., keys) are specified in the request. Alternatively, pass `tag_names=["*"]`
+    to include all tags in the report. Note that tags will be attributed to the entire interval even
+    if they were added or removed at some point within it. If the tag name was not in use during an
+    interval, it will be absent from the tags dictionary in that output row.
+
+    In most cases, billing data will be available in the database that this API queries within
+    minutes, although there may be collection delays. If completeness is important for your use
+    case, we recommend leaving a buffer after the end of the query interval.
+
+    It's also possible to generate reports using the
+    [`modal billing report`](https://modal.com/docs/reference/cli/billing) CLI command. The CLI
+    has a few convenience features for generating reports across relative time ranges.
+
+    """
+    if client is None:
+        client = await _Client.from_env()
+
+    tag_names = tag_names or []
+
+    if end is None:
+        end = datetime.now(timezone.utc)
+
+    if start.tzinfo is None:
+        start = start.replace(tzinfo=timezone.utc)
+    elif start.tzinfo != timezone.utc:
+        raise InvalidError("Timezone-aware 'start' parameter must be in UTC.")
+
+    if end.tzinfo is None:
+        end = end.replace(tzinfo=timezone.utc)
+    elif end.tzinfo != timezone.utc:
+        raise InvalidError("Timezone-aware 'end' parameter must be in UTC.")
+
+    request = api_pb2.WorkspaceBillingReportRequest(
+        resolution=resolution,
+        tag_names=tag_names,
+    )
+    request.start_timestamp.FromDatetime(start)
+    request.end_timestamp.FromDatetime(end)
+
+    rows = []
+    async for pb_item in client.stub.WorkspaceBillingReport.unary_stream(request):
+        item: WorkspaceBillingReportItem = {
+            "object_id": pb_item.object_id,
+            "description": pb_item.description,
+            "environment_name": pb_item.environment_name,
+            "interval_start": pb_item.interval.ToDatetime().replace(tzinfo=timezone.utc),
+            "cost": Decimal(pb_item.cost),
+            "tags": dict(pb_item.tags),
+        }
+        rows.append(item)
+
+    return rows

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_environments.py

@@ -3,7 +3,6 @@ import asyncio
 import builtins
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from typing import Literal
 
 from google.protobuf.empty_pb2 import Empty
@@ -13,7 +12,6 @@ from synchronicity import classproperty
 
 from modal_proto import api_pb2
 
-from ._billing import BillingReportItem
 from ._load_context import LoadContext
 from ._object import _Object
 from ._resolver import Resolver
@@ -320,7 +318,7 @@ class _Environment(_Object, type_prefix="en"):
     # def settings(self) -> EnvironmentSettings:
     #     return self._settings
 
-    def _hydrate_metadata(self, metadata: Message | None):
+    def _hydrate_metadata(self, metadata: Message):
         # Overridden concrete implementation of base class method
         assert metadata and isinstance(metadata, api_pb2.EnvironmentMetadata)
         self._name = metadata.name or None
@@ -394,87 +392,6 @@ class _Environment(_Object, type_prefix="en"):
             client=client,
         )
 
-    @property
-    def billing(self) -> "_EnvironmentBillingManager":
-        return _EnvironmentBillingManager(self)
-
-
-class _EnvironmentBillingManager:
-    """mdmd:namespace
-    Namespace for Environment billing APIs
-    """
-
-    def __init__(self, environment: _Environment):
-        """mdmd:ignore"""
-        self._environment = environment
-
-    async def report(
-        self,
-        *,
-        start: datetime,  # Start of the report, inclusive
-        end: datetime | None = None,  # End of the report, exclusive
-        resolution: str = "d",  # Resolution, e.g. "d" for daily or "h" for hourly
-        tag_names: list[str] | None = None,  # Optional additional metadata to include
-    ) -> list[BillingReportItem]:
-        """Return a cost report for Environment usage, broken down by object and time.
-
-        Args:
-            start: Start of the report, inclusive and rounded to the beginning of the interval.
-                Must be in UTC or timezone-naive (interpreted as UTC).
-            end: End of the report, exclusive. Must be in UTC or timezone-naive. Partial final
-                intervals will be excluded from the report.
-            resolution: Resolution, e.g. "d" for daily or "h" for hourly.
-            tag_names: List of tag names; each row will include the tag name and value in use
-                for that object during the relevant time interval. Pass `["*"]` to include all
-                tags in the report.
-
-        Returns:
-            A list of `BillingReportItem` dataclasses. Each item reports the cost attributed to
-            a specific Modal object during a given time interval. Cost is further broken down by
-            the resource type that generated it (e.g. CPU, Memory, specific GPU usage).
-            Note that the specific resource types included in the breakdown are subject to change
-            as Modal's billing model evolves.
-
-        See also:
-            - [`modal environment billing report`](https://modal.com/docs/cli/latest/environment#modal-environment-billing-report):
-              An environment report CLI that has convenience features around relative time range queries
-              and JSON/CSV output.
-            - [`Workspace.billing.report()`](https://modal.com/docs/sdk/py/latest/modal.Workspace#billingreport):
-              An analogous report API for the entire Workspace.
-
-        """
-        if tag_names is None:
-            tag_names = []
-
-        if end is None:
-            end = datetime.now(timezone.utc)
-
-        if start.tzinfo is None:
-            start = start.replace(tzinfo=timezone.utc)
-        elif start.tzinfo != timezone.utc:
-            raise InvalidError("Timezone-aware 'start' parameter must be in UTC.")
-
-        if end.tzinfo is None:
-            end = end.replace(tzinfo=timezone.utc)
-        elif end.tzinfo != timezone.utc:
-            raise InvalidError("Timezone-aware 'end' parameter must be in UTC.")
-
-        if not self._environment.is_hydrated:
-            await self._environment.hydrate()
-
-        request = api_pb2.WorkspaceBillingReportRequest(
-            resolution=resolution,
-            tag_names=tag_names,
-            environment_ids=[self._environment.object_id],
-        )
-        request.start_timestamp.FromDatetime(start)
-        request.end_timestamp.FromDatetime(end)
-
-        return [
-            BillingReportItem._from_proto(pb_item)
-            async for pb_item in self._environment.client.stub.WorkspaceBillingReport.unary_stream(request)
-        ]
-
 
 ENVIRONMENT_CACHE: dict[str, _Environment] = {}
 

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_functions.py

@@ -661,7 +661,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         min_containers: int | None = None,
         max_containers: int | None = None,
         buffer_containers: int | None = None,
-        scaleup_window: int | None = None,
         scaledown_window: int | None = None,
         max_concurrent_inputs: int | None = None,
         target_concurrent_inputs: int | None = None,
@@ -767,7 +766,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             min_containers=min_containers,
             max_containers=max_containers,
             buffer_containers=buffer_containers,
-            scaleup_window=scaleup_window,
             scaledown_window=scaledown_window,
         )
 
@@ -1159,7 +1157,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         min_containers: int | None = None,
         max_containers: int | None = None,
         buffer_containers: int | None = None,
-        scaleup_window: int | None = None,
         scaledown_window: int | None = None,
         target_concurrency: int | None = None,
     ) -> None:
@@ -1167,7 +1164,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             min_containers=min_containers,
             max_containers=max_containers,
             buffer_containers=buffer_containers,
-            scaleup_window=scaleup_window,
             scaledown_window=scaledown_window,
             target_concurrency=target_concurrency,
         )
@@ -1195,10 +1191,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         its static configuration.
 
         Args:
-            min_containers: Minimum number of containers to keep running.
-            max_containers: Maximum concurrent containers.
-            buffer_containers: Extra containers to keep warm beyond current demand.
-            scaledown_window: Maximum duration (in seconds) idle containers wait before scaling down.
+            min_containers: Minimum number of containers to keep running, or `None` to leave unchanged.
+            max_containers: Maximum concurrent containers, or `None` to leave unchanged.
+            buffer_containers: Extra containers to keep warm beyond demand, or `None` to leave unchanged.
+            scaledown_window: Seconds idle containers wait before scaling down, or `None` to leave unchanged.
 
         Examples:
             ```python notest
@@ -1212,7 +1208,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
             # Extend the scaledown window to increase the amount of time that idle containers stay alive
             f.update_autoscaler(scaledown_window=300)
-            ```
+        ```
 
         """
         # Assert .update_autoscaler() is not called on a method as opposed to the Object. Applicable for Cls only.
@@ -1912,9 +1908,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             **kwargs: Keyword arguments forwarded to the remote function.
 
         Returns:
-            A [`modal.FunctionCall`](https://modal.com/docs/sdk/py/latest/modal.FunctionCall) object
+            A [`modal.FunctionCall`](https://modal.com/docs/reference/modal.FunctionCall) object
             that can later be polled or waited for using
-            [`.get(timeout=...)`](https://modal.com/docs/sdk/py/latest/modal.FunctionCall#get).
+            [`.get(timeout=...)`](https://modal.com/docs/reference/modal.FunctionCall#get).
         """
         self._check_no_web_url("spawn")
         if self._is_generator:
@@ -2065,7 +2061,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         """Returns a structure representing the call graph from a given root
         call ID, along with the status of execution for each node.
 
-        See [`modal.call_graph`](https://modal.com/docs/sdk/py/latest/modal.call_graph) reference page
+        See [`modal.call_graph`](https://modal.com/docs/reference/modal.call_graph) reference page
         for documentation on the structure of the returned `InputInfo` items.
 
         Returns:
@@ -2083,7 +2079,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
         terminate_containers: bool = False,
     ):
         """Cancels the function call, which will stop its execution and mark its inputs as
-        [`TERMINATED`](https://modal.com/docs/sdk/py/latest/modal.call_graph#modalcall_graphinputstatus).
+        [`TERMINATED`](https://modal.com/docs/reference/modal.call_graph#modalcall_graphinputstatus).
 
         If `terminate_containers=True` - the containers running the cancelled inputs are all terminated
         causing any non-cancelled inputs on those containers to be rescheduled in new containers.

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_grpc_client.py

@@ -1,6 +1,6 @@
 # Copyright Modal Labs 2025
 from collections.abc import Collection, Mapping
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Generic, Literal, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, overload
 
 import grpclib.client
 from google.protobuf.message import Message
@@ -178,9 +178,9 @@ class UnaryStreamWrapper(Generic[RequestType, ResponseType]):
 
     async def unary_stream(
         self,
-        request: RequestType,
+        request,
         metadata: Any | None = None,
-    ) -> AsyncGenerator[ResponseType, None]:
+    ):
         from .client import _Client
 
         if self.client._snapshotted:

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_image.py

@@ -9,7 +9,6 @@ import typing
 import warnings
 from collections.abc import Callable, Collection, Sequence
 from dataclasses import dataclass, field
-from functools import wraps
 from inspect import isfunction
 from pathlib import Path, PurePosixPath
 from typing import (
@@ -425,20 +424,6 @@ async def _image_await_build_result(image_id: str, client: _Client) -> api_pb2.I
     return result_response
 
 
-def _requires_image_instance(method):
-    @wraps(method)
-    def wrapper(self, *args, **kwargs):
-        if not isinstance(self, _Image):
-            raise InvalidError(
-                "Image has not been constructed yet. "
-                f"Use one of the static factory methods prior to calling {method.__name__} "
-                "like `modal.Image.debian_slim`, `modal.Image.from_registry`, or `modal.Image.micromamba`"
-            )
-        return method(self, *args, **kwargs)
-
-    return wrapper
-
-
 class _Image(_Object, type_prefix="im"):
     """Base class for container images to run functions in.
 
@@ -770,7 +755,6 @@ class _Image(_Object, type_prefix="im"):
             context_mount_function=lambda: mount,
         )
 
-    @_requires_image_instance
     def add_local_file(self, local_path: str | Path, remote_path: str, *, copy: bool = False) -> "_Image":
         """Adds a local file to the image at `remote_path` within the container.
 
@@ -807,7 +791,6 @@ class _Image(_Object, type_prefix="im"):
         mount = _Mount._from_local_file(local_path, remote_path)
         return self._add_mount_layer_or_copy(mount, copy=copy)
 
-    @_requires_image_instance
     def add_local_dir(
         self,
         local_path: str | Path,
@@ -886,7 +869,6 @@ class _Image(_Object, type_prefix="im"):
         mount = _Mount._add_local_dir(Path(local_path), PurePosixPath(remote_path), ignore=_ignore_fn(ignore))
         return self._add_mount_layer_or_copy(mount, copy=copy)
 
-    @_requires_image_instance
     def add_local_python_source(
         self, *modules: str, copy: bool = False, ignore: Sequence[str] | Callable[[Path], bool] = NON_PYTHON_FILES
     ) -> "_Image":
@@ -1030,7 +1012,6 @@ class _Image(_Object, type_prefix="im"):
             await resolver.load(self, load_context)
         return self
 
-    @_requires_image_instance
     def pip_install(
         self,
         *packages: str | list[str],
@@ -1117,7 +1098,6 @@ class _Image(_Object, type_prefix="im"):
             secrets=secrets,
         )
 
-    @_requires_image_instance
     def pip_install_private_repos(
         self,
         *repositories: str,
@@ -1238,7 +1218,6 @@ class _Image(_Object, type_prefix="im"):
             force_build=self.force_build or force_build,
         )
 
-    @_requires_image_instance
     def pip_install_from_requirements(
         self,
         requirements_txt: str,
@@ -1300,7 +1279,6 @@ class _Image(_Object, type_prefix="im"):
             secrets=secrets,
         )
 
-    @_requires_image_instance
     def pip_install_from_pyproject(
         self,
         pyproject_toml: str,
@@ -1384,7 +1362,6 @@ class _Image(_Object, type_prefix="im"):
             gpu_config=parse_gpu_config(gpu),
         )
 
-    @_requires_image_instance
     def uv_pip_install(
         self,
         *packages: str | list[str],
@@ -1516,7 +1493,6 @@ class _Image(_Object, type_prefix="im"):
             secrets=secrets,
         )
 
-    @_requires_image_instance
     def poetry_install_from_file(
         self,
         poetry_pyproject_toml: str,
@@ -1626,7 +1602,6 @@ class _Image(_Object, type_prefix="im"):
             gpu_config=parse_gpu_config(gpu),
         )
 
-    @_requires_image_instance
     def uv_sync(
         self,
         uv_project_dir: str = "./",
@@ -1810,7 +1785,6 @@ class _Image(_Object, type_prefix="im"):
             gpu_config=parse_gpu_config(gpu),
         )
 
-    @_requires_image_instance
     def dockerfile_commands(
         self,
         *dockerfile_commands: str | list[str],
@@ -1899,7 +1873,6 @@ class _Image(_Object, type_prefix="im"):
             build_args=build_args,
         )
 
-    @_requires_image_instance
     def entrypoint(
         self,
         entrypoint_commands: list[str],
@@ -1920,7 +1893,6 @@ class _Image(_Object, type_prefix="im"):
 
         return self.dockerfile_commands(dockerfile_cmd)
 
-    @_requires_image_instance
     def shell(
         self,
         shell_commands: list[str],
@@ -1941,7 +1913,6 @@ class _Image(_Object, type_prefix="im"):
 
         return self.dockerfile_commands(dockerfile_cmd)
 
-    @_requires_image_instance
     def run_commands(
         self,
         *commands: str | list[str],
@@ -2026,7 +1997,6 @@ class _Image(_Object, type_prefix="im"):
             _namespace=api_pb2.DEPLOYMENT_NAMESPACE_GLOBAL,
         )
 
-    @_requires_image_instance
     def micromamba_install(
         self,
         *packages: str | list[str],
@@ -2558,7 +2528,6 @@ class _Image(_Object, type_prefix="im"):
             _namespace=api_pb2.DEPLOYMENT_NAMESPACE_GLOBAL,
         )
 
-    @_requires_image_instance
     def apt_install(
         self,
         *packages: str | list[str],
@@ -2610,7 +2579,6 @@ class _Image(_Object, type_prefix="im"):
             secrets=secrets,
         )
 
-    @_requires_image_instance
     def run_function(
         self,
         raw_f: Callable[..., Any],
@@ -2729,7 +2697,6 @@ class _Image(_Object, type_prefix="im"):
             force_build=self.force_build or force_build,
         )
 
-    @_requires_image_instance
     def env(self, vars: dict[str, str]) -> "_Image":
         """Sets the environment variables in an Image.
 
@@ -2760,7 +2727,6 @@ class _Image(_Object, type_prefix="im"):
             dockerfile_function=build_dockerfile,
         )
 
-    @_requires_image_instance
     def workdir(self, path: str | PurePosixPath) -> "_Image":
         """Set the working directory for subsequent image build steps and function execution.
 
@@ -2790,7 +2756,6 @@ class _Image(_Object, type_prefix="im"):
             dockerfile_function=build_dockerfile,
         )
 
-    @_requires_image_instance
     def cmd(self, cmd: list[str]) -> "_Image":
         """Set the default command (`CMD`) to run when a container is started.
 
@@ -2819,7 +2784,6 @@ class _Image(_Object, type_prefix="im"):
 
         return self.dockerfile_commands(dockerfile_cmd)
 
-    @_requires_image_instance
     def pipe(
         self,
         func: Callable[Concatenate["modal.Image", P], "modal.Image"],
@@ -2965,16 +2929,17 @@ class _Image(_Object, type_prefix="im"):
         """
         tag = _parse_named_image_ref(name)
 
-        if self._object_id is None:
+        if not self.is_hydrated:
             raise InvalidError("Cannot publish an image that has not been created yet. Call `.build()` first.")
 
-        _client = client or await _Client.from_env()
+        _client = client or self._client
+        assert _client is not None
 
         resolved_env = environment_name or config.get("environment") or ""
 
         await _client.stub.ImagePublish(
             api_pb2.ImagePublishRequest(
-                image_id=self._object_id,
+                image_id=self.object_id,
                 environment_name=resolved_env,
                 is_public=False,
                 tag=tag,

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_partial_function.py

@@ -477,7 +477,7 @@ def _wsgi_app(
     Web Server Gateway Interface (WSGI) is a standard for synchronous Python web apps.
     It has been [succeeded by the ASGI interface](https://asgi.readthedocs.io/en/latest/introduction.html#wsgi-compatibility)
     which is compatible with ASGI and supports additional functionality such as web sockets.
-    Modal supports ASGI via [`asgi_app`](https://modal.com/docs/sdk/py/latest/modal.asgi_app).
+    Modal supports ASGI via [`asgi_app`](https://modal.com/docs/reference/modal.asgi_app).
 
     Examples:
         ```python

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_runtime/gpu_memory_snapshot.py

@@ -17,11 +17,6 @@ from modal.config import config, logger
 
 CUDA_CHECKPOINT_PATH: str = config.get("cuda_checkpoint_path")
 
-# Exit code used to signal to modal-runtime that a GPU memory snapshot restore failed
-# and the task should be retried without a snapshot. This is a sentinel value that must
-# stay in sync with `GPU_SNAPSHOT_RESTORE_FAILED_EXIT_CODE` in runner.rs.
-SNAPSHOT_RESTORE_FAILED_EXIT_CODE: int = 222
-
 # Maximum total duration for each individual `cuda-checkpoint` invocation.
 CUDA_CHECKPOINT_TIMEOUT: float = 3 * 60.0
 

{modal-1.5.1 → modal-1.5.1.dev0}/modal/_runtime/task_lifecycle_manager.py

@@ -180,15 +180,10 @@ class _TaskLifecycleManager:
         if self.function_def._experimental_enable_gpu_snapshot and self.function_def.resources.gpu_config.gpu_type:
             logger.debug("GPU memory snapshot enabled. Attempting to restore GPU memory.")
 
-            try:
-                if self._cuda_checkpoint_session is None:
-                    raise gpu_memory_snapshot.CudaCheckpointException("CudaCheckpointSession not found")
-                self._cuda_checkpoint_session.restore()
-            except gpu_memory_snapshot.CudaCheckpointException as exc:
-                logger.warning(f"GPU memory snapshot restore failed; retrying task without snapshot. Error: {exc}")
-                sys.stderr.flush()
-                # Exit with a sentinel code that the runtime will use to retry the task without a snapshot.
-                os._exit(gpu_memory_snapshot.SNAPSHOT_RESTORE_FAILED_EXIT_CODE)
+            assert self._cuda_checkpoint_session, (
+                "CudaCheckpointSession not found when attempting to restore GPU memory"
+            )
+            self._cuda_checkpoint_session.restore()
 
         self._client = await _Client.from_env()