modal 1.0.3.dev10__tar.gz → 1.0.3.dev12__tar.gz

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.0.3.dev10
+Version: 1.0.3.dev12
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/modal/_runtime/container_io_manager.py

@@ -323,6 +323,7 @@ class _ContainerIOManager:
         self._heartbeat_loop = None
         self._heartbeat_condition = None
         self._waiting_for_memory_snapshot = False
+        self._cuda_checkpoint_session = None
 
         self._is_interactivity_enabled = False
         self._fetching_inputs = True
@@ -881,13 +882,11 @@ class _ContainerIOManager:
         # 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()
+
+            assert self._cuda_checkpoint_session, (
+                "CudaCheckpointSession not found when attempting to restore GPU memory"
+            )
+            self._cuda_checkpoint_session.restore()
 
         # Restore input to default state.
         self.current_input_id = None
@@ -907,14 +906,9 @@ class _ContainerIOManager:
             # 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(
-                        f"Cannot snapshot GPU state if it isn't running. Current GPU state: {gpu_process_state}"
-                    )
 
-                gpu_memory_snapshot.toggle()
-                gpu_memory_snapshot.wait_for_state(gpu_memory_snapshot.CudaCheckpointState.CHECKPOINTED)
+                self._cuda_checkpoint_session = gpu_memory_snapshot.CudaCheckpointSession()
+                self._cuda_checkpoint_session.checkpoint()
 
             # Notify the heartbeat loop that the snapshot phase has begun in order to
             # prevent it from sending heartbeat RPCs

modal-1.0.3.dev12/modal/_runtime/gpu_memory_snapshot.py

@@ -0,0 +1,199 @@
+# 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 subprocess
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+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
+
+
+@dataclass
+class CudaCheckpointProcess:
+    """Contains a reference to a PID with active CUDA session. This also provides
+    methods for checkpointing and restoring GPU memory."""
+
+    pid: int
+    state: CudaCheckpointState
+
+    def toggle(self, target_state: CudaCheckpointState, timeout_secs: float = 5 * 60.0):
+        """Toggle CUDA checkpoint state for current process, moving GPU memory to the
+        CPU and back depending on the current process state when called."""
+        logger.debug(f"PID: {self.pid} Toggling CUDA checkpoint state to {target_state.value}")
+
+        start_time = time.monotonic()
+
+        while self._should_continue_toggle(target_state, start_time, timeout_secs):
+            self._execute_toggle_command()
+            time.sleep(0.1)
+
+        logger.debug(f"PID: {self.pid} Target state {target_state.value} reached")
+
+    def _should_continue_toggle(
+        self, target_state: CudaCheckpointState, start_time: float, timeout_secs: float
+    ) -> bool:
+        """Check if toggle operation should continue based on current state and timeout."""
+        self.refresh_state()
+
+        if self.state == target_state:
+            return False
+
+        if self.state == CudaCheckpointState.FAILED:
+            raise CudaCheckpointException(f"PID: {self.pid} CUDA process state is {self.state}")
+
+        elapsed = time.monotonic() - start_time
+        if elapsed >= timeout_secs:
+            raise CudaCheckpointException(
+                f"PID: {self.pid} Timeout after {elapsed:.2f}s waiting for state {target_state.value}. "
+                f"Current state: {self.state}"
+            )
+
+        return True
+
+    def _execute_toggle_command(self):
+        """Execute the cuda-checkpoint toggle command."""
+        try:
+            subprocess.run(
+                [CUDA_CHECKPOINT_PATH, "--toggle", "--pid", str(self.pid)],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.debug(f"PID: {self.pid} Successfully toggled CUDA checkpoint state")
+        except subprocess.CalledProcessError as e:
+            logger.debug(f"PID: {self.pid} Failed to toggle CUDA checkpoint state: {e.stderr}")
+            raise CudaCheckpointException(e.stderr)
+
+    def refresh_state(self) -> None:
+        """Refreshes the current CUDA checkpoint state for this process."""
+        try:
+            result = subprocess.run(
+                [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(self.pid)],
+                check=True,
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+
+            state_str = result.stdout.strip().lower()
+            self.state = CudaCheckpointState(state_str)
+
+        except subprocess.CalledProcessError as e:
+            logger.debug(f"PID: {self.pid} Failed to get CUDA checkpoint state: {e.stderr}")
+            raise CudaCheckpointException(e.stderr)
+
+
+class CudaCheckpointSession:
+    """Manages the checkpointing state of processes with active CUDA sessions."""
+
+    def __init__(self):
+        self.cuda_processes = self._get_cuda_pids()
+        logger.debug(f"PIDs with CUDA sessions: {[c.pid for c in self.cuda_processes]}")
+
+    def _get_cuda_pids(self) -> list[CudaCheckpointProcess]:
+        """Iterates over all PIDs and identifies the ones that have running
+        CUDA sessions."""
+        cuda_pids: list[CudaCheckpointProcess] = []
+
+        # Get all active process IDs from /proc directory
+        proc_dir = Path("/proc")
+        if not proc_dir.exists():
+            raise CudaCheckpointException(
+                "OS does not have /proc path rendering it incompatible with GPU memory snapshots."
+            )
+
+        for entry in proc_dir.iterdir():
+            if not entry.name.isdigit():
+                continue
+
+            pid = int(entry.name)
+            try:
+                # Call cuda-checkpoint to check if this PID has a CUDA session
+                result = subprocess.run(
+                    [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(pid)],
+                    capture_output=True,
+                    text=True,
+                    timeout=10,
+                )
+
+                # If the command succeeds (return code 0), this PID has a CUDA session
+                if result.returncode == 0:
+                    state_str = result.stdout.strip().lower()
+                    state = CudaCheckpointState(state_str)
+
+                    cuda_checkpoint_process = CudaCheckpointProcess(pid=pid, state=state)
+                    cuda_pids.append(cuda_checkpoint_process)
+
+            # Command failed, which is expected for PIDs without CUDA sessions
+            except subprocess.CalledProcessError:
+                continue
+
+            # Raise other exceptions
+            except subprocess.TimeoutExpired:
+                raise CudaCheckpointException(f"Failed to get CUDA state for PID {pid}")
+            except Exception as e:
+                raise CudaCheckpointException(e)
+
+        # Sort PIDs for ordered checkpointing
+        cuda_pids.sort(key=lambda x: x.pid)
+        return cuda_pids
+
+    def checkpoint(self) -> None:
+        # Validate all states first
+        for proc in self.cuda_processes:
+            if proc.state != CudaCheckpointState.RUNNING:
+                raise CudaCheckpointException(f"CUDA session not in {CudaCheckpointState.RUNNING} state.")
+
+        # Moving state from GPU to CPU can take several seconds per CUDA session.
+        # Make a parallel call per CUDA session.
+        start = time.perf_counter()
+
+        def checkpoint_impl(proc: CudaCheckpointProcess):
+            proc.toggle(CudaCheckpointState.CHECKPOINTED)
+
+        with ThreadPoolExecutor() as executor:
+            list(executor.map(checkpoint_impl, self.cuda_processes))
+
+        elapsed = time.perf_counter() - start
+        logger.debug(f"Checkpointing CUDA sessions took => {elapsed:.3f}s")
+
+    def restore(self) -> None:
+        # Validate all states first
+        for proc in self.cuda_processes:
+            if proc.state != CudaCheckpointState.CHECKPOINTED:
+                raise CudaCheckpointException(f"CUDA session not in {CudaCheckpointState.CHECKPOINTED} state.")
+
+        # See checkpoint() for rationale about parallelism.
+        start = time.perf_counter()
+
+        def restore_process(proc: CudaCheckpointProcess):
+            proc.toggle(CudaCheckpointState.RUNNING)
+
+        with ThreadPoolExecutor() as executor:
+            list(executor.map(restore_process, self.cuda_processes))
+
+        elapsed = time.perf_counter() - start
+        logger.debug(f"Restoring CUDA sessions took => {elapsed:.3f}s")

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/modal/client.pyi

@@ -31,7 +31,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.3.dev10",
+        version: str = "1.0.3.dev12",
     ): ...
     def is_closed(self) -> bool: ...
     @property
@@ -94,7 +94,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.0.3.dev10",
+        version: str = "1.0.3.dev12",
     ): ...
     def is_closed(self) -> bool: ...
     @property

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/modal/functions.pyi

@@ -227,11 +227,11 @@ class Function(
 
     _call_generator: ___call_generator_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
         async def aio(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER: ...
 
-    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -246,12 +246,12 @@ class Function(
         self, *args: modal._functions.P.args, **kwargs: modal._functions.P.kwargs
     ) -> modal._functions.OriginalReturnType: ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
+        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
     ]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER, SUPERSELF]):
@@ -260,11 +260,11 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P, typing_extensions.Self]
 
-    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
         async def aio(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]: ...
 

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/modal.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modal
-Version: 1.0.3.dev10
+Version: 1.0.3.dev12
 Summary: Python client library for Modal
 Author-email: Modal Labs <support@modal.com>
 License: Apache-2.0

{modal-1.0.3.dev10 → modal-1.0.3.dev12}/modal_version/__init__.py

@@ -1,4 +1,4 @@
 # Copyright Modal Labs 2025
 """Supplies the current version of the modal client library."""
 
-__version__ = "1.0.3.dev10"
+__version__ = "1.0.3.dev12"

modal-1.0.3.dev10/modal/_runtime/gpu_memory_snapshot.py

@@ -1,101 +0,0 @@
-# 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:
-        subprocess.run(
-            [
-                CUDA_CHECKPOINT_PATH,
-                "--toggle",
-                "--pid",
-                str(pid),
-            ],
-            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()
-        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")
-            break
-
-        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()