modal 0.73.8__py3-none-any.whl → 0.73.10__py3-none-any.whl

modal/_runtime/container_io_manager.py

@@ -25,6 +25,7 @@ from grpclib import Status
 from synchronicity.async_wrap import asynccontextmanager
 
 import modal_proto.api_pb2
+from modal._runtime import gpu_memory_snapshot
 from modal._serialization import deserialize, serialize, serialize_data_format
 from modal._traceback import extract_traceback, print_exception
 from modal._utils.async_utils import TaskContext, asyncify, synchronize_api, synchronizer
@@ -877,6 +878,17 @@ class _ContainerIOManager:
             if value != "":
                 config.override_locally(key, value)
 
+        # Restore GPU memory.
+        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.")
+            gpu_process_state = gpu_memory_snapshot.get_state()
+            if gpu_process_state != gpu_memory_snapshot.CudaCheckpointState.CHECKPOINTED:
+                raise ValueError(
+                    "Cannot restore GPU state if GPU isn't in a 'checkpointed' state. "
+                    f"Current GPU state: {gpu_process_state}"
+                )
+            gpu_memory_snapshot.toggle()
+
         # Restore input to default state.
         self.current_input_id = None
         self.current_inputs = {}
@@ -892,6 +904,18 @@ class _ContainerIOManager:
 
         # Pause heartbeats since they keep the client connection open which causes the snapshotter to crash
         async with self.heartbeat_condition:
+            # Snapshot GPU memory.
+            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 snapshot GPU memory.")
+                gpu_process_state = gpu_memory_snapshot.get_state()
+                if gpu_process_state != gpu_memory_snapshot.CudaCheckpointState.RUNNING:
+                    raise ValueError(
+                        "Cannot snapshot GPU state if it isn't running. " f"Current GPU state: {gpu_process_state}"
+                    )
+
+                gpu_memory_snapshot.toggle()
+                gpu_memory_snapshot.wait_for_state(gpu_memory_snapshot.CudaCheckpointState.CHECKPOINTED)
+
             # Notify the heartbeat loop that the snapshot phase has begun in order to
             # prevent it from sending heartbeat RPCs
             self._waiting_for_memory_snapshot = True

modal/_runtime/gpu_memory_snapshot.py

@@ -0,0 +1,104 @@
+# Copyright Modal Labs 2022
+#
+# This module provides a simple interface for creating GPU memory snapshots,
+# provising a convenient interface to `cuda-checkpoint` [1]. This is intended
+# to be used in conjunction with memory snapshots.
+#
+# [1] https://github.com/NVIDIA/cuda-checkpoint
+
+import os
+import subprocess
+import time
+from enum import Enum
+
+from modal.config import config, logger
+
+CUDA_CHECKPOINT_PATH: str = config.get("cuda_checkpoint_path")
+
+
+class CudaCheckpointState(Enum):
+    """State representation from the CUDA API: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TYPES.html#group__CUDA__TYPES_1gc96cdda177a2b8c296144567cbea4f23"""
+
+    RUNNING = "running"
+    LOCKED = "locked"
+    CHECKPOINTED = "checkpointed"
+    FAILED = "failed"
+
+
+class CudaCheckpointException(Exception):
+    pass
+
+
+def toggle():
+    """Toggle CUDA checkpoint state for current process, moving GPU memory to the
+    CPU and back depending on the current process state when called."""
+    pid = get_own_pid()
+    logger.debug(f"Toggling CUDA checkpoint state for PID {pid}")
+
+    try:
+        cuda_checkpoint_lock_timeout_ms = 5 * 1000
+        subprocess.run(
+            [
+                CUDA_CHECKPOINT_PATH,
+                "--toggle",
+                "--pid",
+                str(pid),
+                "--timeout",
+                str(cuda_checkpoint_lock_timeout_ms),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        logger.debug("Successfully toggled CUDA checkpoint state")
+
+    except subprocess.CalledProcessError as e:
+        logger.debug(f"Failed to toggle CUDA checkpoint state: {e.stderr}")
+        raise CudaCheckpointException(e.stderr)
+
+
+def get_state() -> CudaCheckpointState:
+    """Get current CUDA checkpoint state for this process."""
+    pid = get_own_pid()
+
+    try:
+        result = subprocess.run(
+            [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(pid)], check=True, capture_output=True, text=True
+        )
+
+        # Parse output to get state
+        state_str = result.stdout.strip().lower()
+        logger.debug(f"Raw state output: {state_str}")
+        return CudaCheckpointState(state_str)
+
+    except subprocess.CalledProcessError as e:
+        logger.debug(f"Failed to get CUDA checkpoint state: {e.stderr}")
+        raise CudaCheckpointException(e.stderr)
+
+
+def wait_for_state(target_state: CudaCheckpointState, timeout_secs: float = 5.0):
+    """Wait for CUDA checkpoint to reach a specific state."""
+    logger.debug(f"Waiting for CUDA checkpoint state {target_state.value}")
+    start_time = time.monotonic()
+
+    while True:
+        current_state = get_state()
+
+        if current_state == target_state:
+            logger.debug(f"Target state {target_state.value} reached")
+
+        if current_state == CudaCheckpointState.FAILED:
+            raise CudaCheckpointException(f"CUDA process state is {current_state}")
+
+        elapsed = time.monotonic() - start_time
+        if elapsed >= timeout_secs:
+            raise CudaCheckpointException(f"Timeout after {elapsed:.2f}s waiting for state {target_state.value}")
+
+        time.sleep(0.1)
+
+
+def get_own_pid():
+    """Returns the Process ID (PID) of the current Python process
+    using only the standard library.
+    """
+    return os.getpid()

modal/cli/run.py

@@ -7,6 +7,7 @@ import re
 import shlex
 import sys
 import time
+from dataclasses import dataclass
 from functools import partial
 from typing import Any, Callable, Optional, get_type_hints
 
@@ -35,6 +36,7 @@ class ParameterMetadata(TypedDict):
     default: Any
     annotation: Any
     type_hint: Any
+    kind: Any
 
 
 class AnyParamType(click.ParamType):
@@ -58,7 +60,13 @@ class NoParserAvailable(InvalidError):
     pass
 
 
-def _get_signature(f: Callable[..., Any], is_method: bool = False) -> dict[str, ParameterMetadata]:
+@dataclass
+class FnSignature:
+    parameters: dict[str, ParameterMetadata]
+    has_variadic_args: bool
+
+
+def _get_signature(f: Callable[..., Any], is_method: bool = False) -> FnSignature:
     try:
         type_hints = get_type_hints(f)
     except Exception as exc:
@@ -66,18 +74,29 @@ def _get_signature(f: Callable[..., Any], is_method: bool = False) -> dict[str,
         msg = "Unable to generate command line interface for app entrypoint. See traceback above for details."
         raise ExecutionError(msg) from exc
 
+    has_variadic_args = False
+
     if is_method:
         self = None  # Dummy, doesn't matter
         f = functools.partial(f, self)
+
     signature: dict[str, ParameterMetadata] = {}
     for param in inspect.signature(f).parameters.values():
-        signature[param.name] = {
-            "name": param.name,
-            "default": param.default,
-            "annotation": param.annotation,
-            "type_hint": type_hints.get(param.name, "Any"),
-        }
-    return signature
+        if param.kind == inspect.Parameter.VAR_POSITIONAL:
+            has_variadic_args = True
+        else:
+            signature[param.name] = {
+                "name": param.name,
+                "default": param.default,
+                "annotation": param.annotation,
+                "type_hint": type_hints.get(param.name, "Any"),
+                "kind": param.kind,
+            }
+
+    if has_variadic_args and len(signature) > 0:
+        raise InvalidError("Functions with variable-length positional arguments (*args) cannot have other parameters.")
+
+    return FnSignature(signature, has_variadic_args)
 
 
 def _get_param_type_as_str(annot: Any) -> str:
@@ -96,17 +115,18 @@ def _get_param_type_as_str(annot: Any) -> str:
     return annot_str
 
 
-def _add_click_options(func, signature: dict[str, ParameterMetadata]):
+def _add_click_options(func, parameters: dict[str, ParameterMetadata]):
     """Adds @click.option based on function signature
 
     Kind of like typer, but using options instead of positional arguments
     """
-    for param in signature.values():
+    for param in parameters.values():
         param_type_str = _get_param_type_as_str(param["type_hint"])
         param_name = param["name"].replace("_", "-")
         cli_name = "--" + param_name
         if param_type_str == "bool":
             cli_name += "/--no-" + param_name
+
         parser = option_parsers.get(param_type_str)
         if parser is None:
             msg = f"Parameter `{param_name}` has unparseable annotation: {param['annotation']!r}"
@@ -145,9 +165,15 @@ def _write_local_result(result_path: str, res: Any):
         fid.write(res)
 
 
-def _make_click_function(app, inner: Callable[[dict[str, Any]], Any]):
+def _make_click_function(app, signature: FnSignature, inner: Callable[[tuple[str, ...], dict[str, Any]], Any]):
     @click.pass_context
     def f(ctx, **kwargs):
+        if signature.has_variadic_args:
+            assert len(kwargs) == 0
+            args = ctx.args
+        else:
+            args = ()
+
         show_progress: bool = ctx.obj["show_progress"]
         with enable_output(show_progress):
             with run_app(
@@ -156,7 +182,7 @@ def _make_click_function(app, inner: Callable[[dict[str, Any]], Any]):
                 environment_name=ctx.obj["env"],
                 interactive=ctx.obj["interactive"],
             ):
-                res = inner(kwargs)
+                res = inner(args, kwargs)
 
             if result_path := ctx.obj["result_path"]:
                 _write_local_result(result_path, res)
@@ -168,23 +194,33 @@ def _get_click_command_for_function(app: App, function: Function):
     if function.is_generator:
         raise InvalidError("`modal run` is not supported for generator functions")
 
-    signature: dict[str, ParameterMetadata] = _get_signature(function.info.raw_f)
+    signature = _get_signature(function.info.raw_f)
+
+    def _inner(args, click_kwargs):
+        return function.remote(*args, **click_kwargs)
 
-    def _inner(click_kwargs):
-        return function.remote(**click_kwargs)
+    f = _make_click_function(app, signature, _inner)
 
-    f = _make_click_function(app, _inner)
+    with_click_options = _add_click_options(f, signature.parameters)
 
-    with_click_options = _add_click_options(f, signature)
-    return click.command(with_click_options)
+    if signature.has_variadic_args:
+        return click.command(context_settings={"ignore_unknown_options": True, "allow_extra_args": True})(
+            with_click_options
+        )
+    else:
+        return click.command(with_click_options)
 
 
 def _get_click_command_for_cls(app: App, method_ref: MethodReference):
-    signature: dict[str, ParameterMetadata]
+    parameters: dict[str, ParameterMetadata]
     cls = method_ref.cls
     method_name = method_ref.method_name
 
     cls_signature = _get_signature(cls._get_user_cls())
+
+    if cls_signature.has_variadic_args:
+        raise InvalidError("Modal classes cannot have variable-length positional arguments (*args).")
+
     partial_functions = cls._get_partial_functions()
 
     if method_name in ("*", ""):
@@ -202,27 +238,35 @@ def _get_click_command_for_cls(app: App, method_ref: MethodReference):
     fun_signature = _get_signature(partial_function._get_raw_f(), is_method=True)
 
     # TODO(erikbern): assert there's no overlap?
-    signature = dict(**cls_signature, **fun_signature)  # Pool all arguments
+    parameters = dict(**cls_signature.parameters, **fun_signature.parameters)  # Pool all arguments
 
-    def _inner(click_kwargs):
+    def _inner(args, click_kwargs):
         # unpool class and method arguments
         # TODO(erikbern): this code is a bit hacky
-        cls_kwargs = {k: click_kwargs[k] for k in cls_signature}
-        fun_kwargs = {k: click_kwargs[k] for k in fun_signature}
+        cls_kwargs = {k: click_kwargs[k] for k in cls_signature.parameters}
+        fun_kwargs = {k: click_kwargs[k] for k in fun_signature.parameters}
 
         instance = cls(**cls_kwargs)
         method: Function = getattr(instance, method_name)
-        return method.remote(**fun_kwargs)
+        return method.remote(*args, **fun_kwargs)
+
+    f = _make_click_function(app, fun_signature, _inner)
+    with_click_options = _add_click_options(f, parameters)
 
-    f = _make_click_function(app, _inner)
-    with_click_options = _add_click_options(f, signature)
-    return click.command(with_click_options)
+    if fun_signature.has_variadic_args:
+        return click.command(context_settings={"ignore_unknown_options": True, "allow_extra_args": True})(
+            with_click_options
+        )
+    else:
+        return click.command(with_click_options)
 
 
 def _get_click_command_for_local_entrypoint(app: App, entrypoint: LocalEntrypoint):
     func = entrypoint.info.raw_f
     isasync = inspect.iscoroutinefunction(func)
 
+    signature = _get_signature(func)
+
     @click.pass_context
     def f(ctx, *args, **kwargs):
         if ctx.obj["detach"]:
@@ -231,6 +275,10 @@ def _get_click_command_for_local_entrypoint(app: App, entrypoint: LocalEntrypoin
                 "triggered Modal function alive after the parent process has been killed or disconnected."
             )
 
+        if signature.has_variadic_args:
+            assert len(args) == 0 and len(kwargs) == 0
+            args = ctx.args
+
         show_progress: bool = ctx.obj["show_progress"]
         with enable_output(show_progress):
             with run_app(
@@ -250,8 +298,14 @@ def _get_click_command_for_local_entrypoint(app: App, entrypoint: LocalEntrypoin
             if result_path := ctx.obj["result_path"]:
                 _write_local_result(result_path, res)
 
-    with_click_options = _add_click_options(f, _get_signature(func))
-    return click.command(with_click_options)
+    with_click_options = _add_click_options(f, signature.parameters)
+
+    if signature.has_variadic_args:
+        return click.command(context_settings={"ignore_unknown_options": True, "allow_extra_args": True})(
+            with_click_options
+        )
+    else:
+        return click.command(with_click_options)
 
 
 def _get_runnable_list(all_usable_commands: list[CLICommand]) -> str:

modal/client.pyi

@@ -27,7 +27,7 @@ class _Client:
     _snapshotted: bool
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.73.8"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.73.10"
     ): ...
     def is_closed(self) -> bool: ...
     @property
@@ -85,7 +85,7 @@ class Client:
     _snapshotted: bool
 
     def __init__(
-        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.73.8"
+        self, server_url: str, client_type: int, credentials: typing.Optional[tuple[str, str]], version: str = "0.73.10"
     ): ...
     def is_closed(self) -> bool: ...
     @property

modal/config.py

@@ -223,6 +223,7 @@ _SETTINGS = {
     "strict_parameters": _Setting(False, transform=_to_boolean),  # For internal/experimental use
     "snapshot_debug": _Setting(False, transform=_to_boolean),
     "client_retries": _Setting(False, transform=_to_boolean),  # For internal testing.
+    "cuda_checkpoint_path": _Setting("/__modal/.bin/cuda-checkpoint"),  # Used for snapshotting GPU memory.
 }
 
 

{modal-0.73.8.dist-info → modal-0.73.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: modal
-Version: 0.73.8
+Version: 0.73.10
 Summary: Python client library for Modal
 Author: Modal Labs
 Author-email: support@modal.com

{modal-0.73.8.dist-info → modal-0.73.10.dist-info}/RECORD

@@ -21,12 +21,12 @@ modal/app.py,sha256=wRygVSrWH8iIqhDAAl2Ww_RAkz8MCJZ0Jt9qYZCF6SA,44626
 modal/app.pyi,sha256=lxiuWzE_OLb3WHg-H7Pek9DGBuCUzZ55P594VhJL5LA,26113
 modal/call_graph.py,sha256=1g2DGcMIJvRy-xKicuf63IVE98gJSnQsr8R_NVMptNc,2581
 modal/client.py,sha256=8SQawr7P1PNUCq1UmJMUQXG2jIo4Nmdcs311XqrNLRE,15276
-modal/client.pyi,sha256=kOt17KK6b51zHs9T5vme0cCrLzJu-nw1DTE1kSuxXcQ,7591
+modal/client.pyi,sha256=zquF7hra8lDrBysJUafRuLtmDHnUXpEU1nt9wRDDCUs,7593
 modal/cloud_bucket_mount.py,sha256=YOe9nnvSr4ZbeCn587d7_VhE9IioZYRvF9VYQTQux08,5914
 modal/cloud_bucket_mount.pyi,sha256=30T3K1a89l6wzmEJ_J9iWv9SknoGqaZDx59Xs-ZQcmk,1607
 modal/cls.py,sha256=kNnZrBYVXOhgEXU0rDWk2Hr-bQRrsZkMKDgC-TD_6Bs,31063
 modal/cls.pyi,sha256=gb6QNwfX3HSJfcZXPY36N9ywF7aBJTwwtoARnf3G1HQ,8877
-modal/config.py,sha256=BzhZYUUwOmvVwf6x5kf0ywMC257s648dmuhsnB6g3gk,11041
+modal/config.py,sha256=XT1W4Y9PVkbYMAXjJRshvQEPDhZmnfW_ZRMwl8XKoqA,11149
 modal/container_process.py,sha256=WTqLn01dJPVkPpwR_0w_JH96ceN5mV4TGtiu1ZR2RRA,6108
 modal/container_process.pyi,sha256=Hf0J5JyDdCCXBJSKx6gvkPOo0XrztCm78xzxamtzUjQ,2828
 modal/dict.py,sha256=vc5lQVqzeDUCb4fRjnOlqYK2GmBb0fIhZmvB0xIBG0U,12921
@@ -83,8 +83,9 @@ modal/volume.py,sha256=JAWeDvoAG95tMBv-fYIERyHsJPS_X_xGpxRRmYtb6j0,30096
 modal/volume.pyi,sha256=kTsXarphjZILXci84LQy7EyC84eXUs5-7D62IM5q3eE,12491
 modal/_runtime/__init__.py,sha256=MIEP8jhXUeGq_eCjYFcqN5b1bxBM4fdk0VESpjWR0fc,28
 modal/_runtime/asgi.py,sha256=c4hmaMW1pLo-cm7ouriJjieuFm4ZF6D2LMy0638sfOs,22139
-modal/_runtime/container_io_manager.py,sha256=QVWMCvJatd2696wsauzXl20psxCYsR0d_CHeS5ceTsU,43201
+modal/_runtime/container_io_manager.py,sha256=L6qv-Mo3mN3ttR5GX-G36cUhH_oz8wdP5WG0HT5FFzg,44619
 modal/_runtime/execution_context.py,sha256=E6ofm6j1POXGPxS841X3V7JU6NheVb8OkQc7JpLq4Kg,2712
+modal/_runtime/gpu_memory_snapshot.py,sha256=vV6igsqN9CxOoH91kUkuaZQ32QfX5wdoXIS-6MIYX2Y,3315
 modal/_runtime/telemetry.py,sha256=T1RoAGyjBDr1swiM6pPsGRSITm7LI5FDK18oNXxY08U,5163
 modal/_runtime/user_code_imports.py,sha256=zl_Mq9dsrVF62x3w-iNK1YAhZKYAXeFaGpd4G7AySTc,14746
 modal/_utils/__init__.py,sha256=waLjl5c6IPDhSsdWAm9Bji4e2PVxamYABKAze6CHVXY,28
@@ -124,7 +125,7 @@ modal/cli/launch.py,sha256=pzQt2QlcrbIUU0MVzWWPAvMQ6MCyqsHZ0X9JcV-sY04,3242
 modal/cli/network_file_system.py,sha256=eq3JnwjbfFNsJodIyANHL06ByYc3BSavzdmu8C96cHA,7948
 modal/cli/profile.py,sha256=rLXfjJObfPNjaZvNfHGIKqs7y9bGYyGe-K7V0w-Ni0M,3110
 modal/cli/queues.py,sha256=6gTu76dzBtPN5eQVsLrvQpuru5jI9ZCWK5Eh8J8XhaM,4498
-modal/cli/run.py,sha256=_hYvRwBkrbakN3oAuoiaMdYytBtN4BoryrYxPAVIIvc,20065
+modal/cli/run.py,sha256=hXCvqQM6w7dXz0254q5Sr0UNR5JC6cEA6fBdrcshCcg,21784
 modal/cli/secret.py,sha256=uQpwYrMY98iMCWeZOQTcktOYhPTZ8IHnyealDc2CZqo,4206
 modal/cli/token.py,sha256=mxSgOWakXG6N71hQb1ko61XAR9ZGkTMZD-Txn7gmTac,1924
 modal/cli/utils.py,sha256=hZmjyzcPjDnQSkLvycZD2LhGdcsfdZshs_rOU78EpvI,3717
@@ -170,10 +171,10 @@ modal_proto/options_pb2_grpc.pyi,sha256=CImmhxHsYnF09iENPoe8S4J-n93jtgUYD2JPAc0y
 modal_proto/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modal_version/__init__.py,sha256=wiJQ53c-OMs0Xf1UeXOxQ7FwlV1VzIjnX6o-pRYZ_Pk,470
 modal_version/__main__.py,sha256=2FO0yYQQwDTh6udt1h-cBnGd1c4ZyHnHSI4BksxzVac,105
-modal_version/_version_generated.py,sha256=a8LMamoFRZKpp6SoTp5bDH1HRWGljeWLZesEIHf3WIE,148
-modal-0.73.8.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-modal-0.73.8.dist-info/METADATA,sha256=wsn98lBqsvSVbab7C3DqlgtiyUF8vmevqqGNf7Ir6WE,2329
-modal-0.73.8.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-modal-0.73.8.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
-modal-0.73.8.dist-info/top_level.txt,sha256=1nvYbOSIKcmU50fNrpnQnrrOpj269ei3LzgB6j9xGqg,64
-modal-0.73.8.dist-info/RECORD,,
+modal_version/_version_generated.py,sha256=bRP1kitLENe-C1R8SFEB0RNXUvVdmynlJRPg8FgsxXA,149
+modal-0.73.10.dist-info/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+modal-0.73.10.dist-info/METADATA,sha256=o2yy7RC_d68S0vZZyIem-3gJG0H7EDGGXtmHHcPzmEE,2330
+modal-0.73.10.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+modal-0.73.10.dist-info/entry_points.txt,sha256=An-wYgeEUnm6xzrAP9_NTSTSciYvvEWsMZILtYrvpAI,46
+modal-0.73.10.dist-info/top_level.txt,sha256=1nvYbOSIKcmU50fNrpnQnrrOpj269ei3LzgB6j9xGqg,64
+modal-0.73.10.dist-info/RECORD,,

modal_version/_version_generated.py

@@ -1,4 +1,4 @@
 # Copyright Modal Labs 2025
 
 # Note: Reset this value to -1 whenever you make a minor `0.X` release of the client.
-build_number = 8  # git: 6a2791f
+build_number = 10  # git: 6951dc1