modal 1.0.3.dev10__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_runtime/execution_context.pyi

@@ -3,22 +3,79 @@ import contextvars
 import typing
 import typing_extensions
 
-def is_local() -> bool: ...
-async def _interact() -> None: ...
+def is_local() -> bool:
+    """Returns if we are currently on the machine launching/deploying a Modal app
+
+    Returns `True` when executed locally on the user's machine.
+    Returns `False` when executed from a Modal container in the cloud.
+    """
+    ...
+
+async def _interact() -> None:
+    """Enable interactivity with user input inside a Modal container.
+
+    See the [interactivity guide](https://modal.com/docs/guide/developing-debugging#interactivity)
+    for more information on how to use this function.
+    """
+    ...
 
 class __interact_spec(typing_extensions.Protocol):
-    def __call__(self, /) -> None: ...
-    async def aio(self, /) -> None: ...
+    def __call__(self, /) -> None:
+        """Enable interactivity with user input inside a Modal container.
+
+        See the [interactivity guide](https://modal.com/docs/guide/developing-debugging#interactivity)
+        for more information on how to use this function.
+        """
+        ...
+
+    async def aio(self, /) -> None:
+        """Enable interactivity with user input inside a Modal container.
+
+        See the [interactivity guide](https://modal.com/docs/guide/developing-debugging#interactivity)
+        for more information on how to use this function.
+        """
+        ...
 
 interact: __interact_spec
 
-def current_input_id() -> typing.Optional[str]: ...
-def current_function_call_id() -> typing.Optional[str]: ...
+def current_input_id() -> typing.Optional[str]:
+    """Returns the input ID for the current input.
+
+    Can only be called from Modal function (i.e. in a container context).
+
+    ```python
+    from modal import current_input_id
+
+    @app.function()
+    def process_stuff():
+        print(f"Starting to process {current_input_id()}")
+    ```
+    """
+    ...
+
+def current_function_call_id() -> typing.Optional[str]:
+    """Returns the function call ID for the current input.
+
+    Can only be called from Modal function (i.e. in a container context).
+
+    ```python
+    from modal import current_function_call_id
+
+    @app.function()
+    def process_stuff():
+        print(f"Starting to process input from {current_function_call_id()}")
+    ```
+    """
+    ...
+
+def current_attempt_token() -> typing.Optional[str]: ...
 def _set_current_context_ids(
-    input_ids: list[str], function_call_ids: list[str]
+    input_ids: list[str], function_call_ids: list[str], attempt_tokens: list[str]
 ) -> collections.abc.Callable[[], None]: ...
 def _import_context(): ...
 
 _current_input_id: contextvars.ContextVar
 
 _current_function_call_id: contextvars.ContextVar
+
+_current_attempt_token: contextvars.ContextVar

modal/_runtime/gpu_memory_snapshot.py

@@ -1,23 +1,34 @@
 # 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
+# providing 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 concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
+from typing import List, Optional
 
 from modal.config import config, logger
 
 CUDA_CHECKPOINT_PATH: str = config.get("cuda_checkpoint_path")
 
+# Maximum total duration for an entire toggle operation.
+CUDA_CHECKPOINT_TOGGLE_TIMEOUT: float = 5 * 60.0
+
+# Maximum total duration for each individual `cuda-checkpoint` invocation.
+CUDA_CHECKPOINT_TIMEOUT: float = 90
+
 
 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"""
+    """State representation from the CUDA API [1].
+
+    [1] https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TYPES.html"""
 
     RUNNING = "running"
     LOCKED = "locked"
@@ -26,76 +37,270 @@ class CudaCheckpointState(Enum):
 
 
 class CudaCheckpointException(Exception):
+    """Exception raised for CUDA checkpoint operations."""
+
     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}")
+@dataclass
+class CudaCheckpointProcess:
+    """Contains a reference to a PID with active CUDA session. This also provides
+    methods for checkpointing and restoring GPU memory."""
 
-    try:
-        subprocess.run(
-            [
-                CUDA_CHECKPOINT_PATH,
-                "--toggle",
-                "--pid",
-                str(pid),
-            ],
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-        logger.debug("Successfully toggled CUDA checkpoint state")
+    pid: int
+    state: CudaCheckpointState
 
-    except subprocess.CalledProcessError as e:
-        logger.debug(f"Failed to toggle CUDA checkpoint state: {e.stderr}")
-        raise CudaCheckpointException(e.stderr)
+    def toggle(self, target_state: CudaCheckpointState, skip_first_refresh: bool = False) -> None:
+        """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()
+        retry_count = 0
+        max_retries = 3
 
-def get_state() -> CudaCheckpointState:
-    """Get current CUDA checkpoint state for this process."""
-    pid = get_own_pid()
+        attempts = 0
+        while self._should_continue_toggle(
+            target_state, start_time, refresh=not (skip_first_refresh and attempts == 0)
+        ):
+            attempts += 1
+            try:
+                self._execute_toggle_command()
+                # Use exponential backoff for retries
+                sleep_time = min(0.1 * (2**retry_count), 1.0)
+                time.sleep(sleep_time)
+                retry_count = 0
+            except CudaCheckpointException as e:
+                retry_count += 1
+                if retry_count >= max_retries:
+                    raise CudaCheckpointException(
+                        f"PID: {self.pid} Failed to toggle state after {max_retries} retries: {e}"
+                    )
+                logger.debug(f"PID: {self.pid} Retry {retry_count}/{max_retries} after error: {e}")
+                time.sleep(0.5 * retry_count)
 
-    try:
-        result = subprocess.run(
-            [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(pid)], check=True, capture_output=True, text=True
-        )
+        logger.debug(f"PID: {self.pid} Target state {target_state.value} reached")
 
-        # Parse output to get state
-        state_str = result.stdout.strip().lower()
-        return CudaCheckpointState(state_str)
+    def _should_continue_toggle(
+        self, target_state: CudaCheckpointState, start_time: float, refresh: bool = True
+    ) -> bool:
+        """Check if toggle operation should continue based on current state and timeout."""
+        if refresh:
+            self.refresh_state()
 
-    except subprocess.CalledProcessError as e:
-        logger.debug(f"Failed to get CUDA checkpoint state: {e.stderr}")
-        raise CudaCheckpointException(e.stderr)
+        if self.state == target_state:
+            return False
 
+        if self.state == CudaCheckpointState.FAILED:
+            raise CudaCheckpointException(f"PID: {self.pid} CUDA process state is {self.state}")
 
-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()
+        elapsed = time.monotonic() - start_time
+        if elapsed >= CUDA_CHECKPOINT_TOGGLE_TIMEOUT:
+            raise CudaCheckpointException(
+                f"PID: {self.pid} Timeout after {elapsed:.2f}s waiting for state {target_state.value}. "
+                f"Current state: {self.state}"
+            )
 
-    while True:
-        current_state = get_state()
+        return True
 
-        if current_state == target_state:
-            logger.debug(f"Target state {target_state.value} reached")
-            break
+    def _execute_toggle_command(self) -> None:
+        """Execute the cuda-checkpoint toggle command."""
+        try:
+            _ = subprocess.run(
+                [CUDA_CHECKPOINT_PATH, "--toggle", "--pid", str(self.pid)],
+                check=True,
+                capture_output=True,
+                text=True,
+                timeout=CUDA_CHECKPOINT_TIMEOUT,
+            )
+            logger.debug(f"PID: {self.pid} Successfully toggled CUDA checkpoint state")
+        except subprocess.CalledProcessError as e:
+            error_msg = f"PID: {self.pid} Failed to toggle CUDA checkpoint state: {e.stderr}"
+            logger.debug(error_msg)
+            raise CudaCheckpointException(error_msg)
+        except subprocess.TimeoutExpired:
+            error_msg = f"PID: {self.pid} Toggle command timed out"
+            logger.debug(error_msg)
+            raise CudaCheckpointException(error_msg)
 
-        if current_state == CudaCheckpointState.FAILED:
-            raise CudaCheckpointException(f"CUDA process state is {current_state}")
+    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=CUDA_CHECKPOINT_TIMEOUT,
+            )
 
-        elapsed = time.monotonic() - start_time
-        if elapsed >= timeout_secs:
-            raise CudaCheckpointException(f"Timeout after {elapsed:.2f}s waiting for state {target_state.value}")
+            state_str = result.stdout.strip().lower()
+            self.state = CudaCheckpointState(state_str)
+
+        except subprocess.CalledProcessError as e:
+            error_msg = f"PID: {self.pid} Failed to get CUDA checkpoint state: {e.stderr}"
+            logger.debug(error_msg)
+            raise CudaCheckpointException(error_msg)
+        except subprocess.TimeoutExpired:
+            error_msg = f"PID: {self.pid} Get state command timed out"
+            logger.debug(error_msg)
+            raise CudaCheckpointException(error_msg)
+
+
+class CudaCheckpointSession:
+    """Manages the checkpointing state of processes with active CUDA sessions."""
+
+    def __init__(self):
+        self.cuda_processes = self._get_cuda_pids()
+        if self.cuda_processes:
+            logger.debug(
+                f"Found {len(self.cuda_processes)} PID(s) with CUDA sessions: {[c.pid for c in self.cuda_processes]}"
+            )
+        else:
+            logger.debug("No CUDA sessions found.")
+
+    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."
+            )
+
+        # Get all numeric directories (PIDs) from /proc
+        pid_dirs = [entry for entry in proc_dir.iterdir() if entry.name.isdigit()]
+
+        # Use ThreadPoolExecutor to check PIDs in parallel for better performance
+        with ThreadPoolExecutor(max_workers=min(50, len(pid_dirs))) as executor:
+            future_to_pid = {
+                executor.submit(self._check_cuda_session, int(entry.name)): int(entry.name) for entry in pid_dirs
+            }
+
+            for future in as_completed(future_to_pid):
+                pid = future_to_pid[future]
+                try:
+                    cuda_process = future.result()
+                    if cuda_process:
+                        cuda_pids.append(cuda_process)
+                except Exception as e:
+                    logger.debug(f"Error checking PID {pid}: {e}")
+
+        # Sort PIDs for ordered checkpointing
+        cuda_pids.sort(key=lambda x: x.pid)
+        return cuda_pids
+
+    def _check_cuda_session(self, pid: int) -> Optional[CudaCheckpointProcess]:
+        """Check if a specific PID has a CUDA session."""
+        try:
+            result = subprocess.run(
+                [CUDA_CHECKPOINT_PATH, "--get-state", "--pid", str(pid)],
+                capture_output=True,
+                text=True,
+                # This should be quick since no checkpoint has taken place yet
+                timeout=5,
+            )
+
+            # 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)
+                return CudaCheckpointProcess(pid=pid, state=state)
+
+        except subprocess.CalledProcessError:
+            # Command failed, which is expected for PIDs without CUDA sessions
+            pass
+        except subprocess.TimeoutExpired:
+            logger.debug(f"Timeout checking CUDA state for PID {pid}")
+        except Exception as e:
+            logger.debug(f"Error checking PID {pid}: {e}")
+
+        return None
+
+    def checkpoint(self) -> None:
+        """Checkpoint all CUDA processes, moving GPU memory to CPU."""
+        if not self.cuda_processes:
+            logger.debug("No CUDA processes to checkpoint.")
+            return
+
+        # Validate all states first
+        for proc in self.cuda_processes:
+            proc.refresh_state()  # Refresh state before validation
+            if proc.state != CudaCheckpointState.RUNNING:
+                raise CudaCheckpointException(
+                    f"PID {proc.pid}: CUDA session not in {CudaCheckpointState.RUNNING.value} state. "
+                    f"Current state: {proc.state.value}"
+                )
+
+        # 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) -> None:
+            proc.toggle(CudaCheckpointState.CHECKPOINTED)
+
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(checkpoint_impl, proc) for proc in self.cuda_processes]
+
+            # Wait for all futures and collect any exceptions
+            exceptions = []
+            for future in as_completed(futures):
+                try:
+                    future.result()
+                except Exception as e:
+                    exceptions.append(e)
+
+            if exceptions:
+                raise CudaCheckpointException(
+                    f"Failed to checkpoint {len(exceptions)} processes: {'; '.join(str(e) for e in exceptions)}"
+                )
+
+        elapsed = time.perf_counter() - start
+        logger.debug(f"Checkpointing {len(self.cuda_processes)} CUDA sessions took => {elapsed:.3f}s")
+
+    def restore(self) -> None:
+        """Restore all CUDA processes, moving memory back from CPU to GPU."""
+        if not self.cuda_processes:
+            logger.debug("No CUDA sessions to restore.")
+            return
+
+        # See checkpoint() for rationale about parallelism.
+        start = time.perf_counter()
+
+        def restore_process(proc: CudaCheckpointProcess) -> None:
+            proc.toggle(CudaCheckpointState.RUNNING, skip_first_refresh=True)
+
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(restore_process, proc) for proc in self.cuda_processes]
+
+            # Wait for all futures and collect any exceptions
+            exceptions = []
+            for future in as_completed(futures):
+                try:
+                    future.result()
+                except Exception as e:
+                    exceptions.append(e)
+
+            if exceptions:
+                raise CudaCheckpointException(
+                    f"Failed to restore {len(exceptions)} processes: {'; '.join(str(e) for e in exceptions)}"
+                )
 
-        time.sleep(0.1)
+        elapsed = time.perf_counter() - start
+        logger.debug(f"Restoring {len(self.cuda_processes)} CUDA session(s) took => {elapsed:.3f}s")
 
+    def get_process_count(self) -> int:
+        """Get the number of CUDA processes managed by this session."""
+        return len(self.cuda_processes)
 
-def get_own_pid():
-    """Returns the Process ID (PID) of the current Python process
-    using only the standard library.
-    """
-    return os.getpid()
+    def get_process_states(self) -> List[tuple[int, CudaCheckpointState]]:
+        """Get current states of all managed processes."""
+        states = []
+        for proc in self.cuda_processes:
+            proc.refresh_state()
+            states.append((proc.pid, proc.state))
+        return states

modal/_runtime/user_code_imports.py

@@ -29,7 +29,7 @@ class FinalizedFunction:
     callable: Callable[..., Any]
     is_async: bool
     is_generator: bool
-    data_format: int  # api_pb2.DataFormat
+    supported_output_formats: Sequence["api_pb2.DataFormat.ValueType"]
     lifespan_manager: Optional["LifespanManager"] = None
 
 
@@ -93,9 +93,9 @@ def construct_webhook_callable(
 
 @dataclass
 class ImportedFunction(Service):
-    user_cls_instance: Any
     app: modal.app._App
     service_deps: Optional[Sequence["modal._object._Object"]]
+    user_cls_instance = None
 
     _user_defined_callable: Callable[..., Any]
 
@@ -108,6 +108,7 @@ class ImportedFunction(Service):
         is_generator = fun_def.function_type == api_pb2.Function.FUNCTION_TYPE_GENERATOR
 
         webhook_config = fun_def.webhook_config
+
         if not webhook_config.type:
             # for non-webhooks, the runnable is straight forward:
             return {
@@ -115,7 +116,10 @@ class ImportedFunction(Service):
                     callable=self._user_defined_callable,
                     is_async=is_async,
                     is_generator=is_generator,
-                    data_format=api_pb2.DATA_FORMAT_PICKLE,
+                    supported_output_formats=fun_def.supported_output_formats
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]` is only
+                    # needed for tests
+                    or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR],
                 )
             }
 
@@ -129,7 +133,8 @@ class ImportedFunction(Service):
                 lifespan_manager=lifespan_manager,
                 is_async=True,
                 is_generator=True,
-                data_format=api_pb2.DATA_FORMAT_ASGI,
+                # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_ASGI]` is only needed for tests
+                supported_output_formats=fun_def.supported_output_formats or [api_pb2.DATA_FORMAT_ASGI],
             )
         }
 
@@ -154,6 +159,7 @@ class ImportedClass(Service):
             # Use the function definition for whether this is a generator (overriden by webhooks)
             is_generator = _partial.params.is_generator
             webhook_config = _partial.params.webhook_config
+            method_def = fun_def.method_definitions[method_name]
 
             bound_func = user_func.__get__(self.user_cls_instance)
 
@@ -163,7 +169,10 @@ class ImportedClass(Service):
                     callable=bound_func,
                     is_async=is_async,
                     is_generator=bool(is_generator),
-                    data_format=api_pb2.DATA_FORMAT_PICKLE,
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]` is only
+                    # needed for tests
+                    supported_output_formats=method_def.supported_output_formats
+                    or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR],
                 )
             else:
                 web_callable, lifespan_manager = construct_webhook_callable(
@@ -174,7 +183,8 @@ class ImportedClass(Service):
                     lifespan_manager=lifespan_manager,
                     is_async=True,
                     is_generator=True,
-                    data_format=api_pb2.DATA_FORMAT_ASGI,
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_ASGI]` is only needed for tests
+                    supported_output_formats=method_def.supported_output_formats or [api_pb2.DATA_FORMAT_ASGI],
                 )
             finalized_functions[method_name] = finalized_function
         return finalized_functions
@@ -199,7 +209,6 @@ def get_user_class_instance(_cls: modal.cls._Cls, args: tuple[Any, ...], kwargs:
 
 def import_single_function_service(
     function_def: api_pb2.Function,
-    ser_cls: Optional[type],  # used only for @build functions
     ser_fun: Optional[Callable[..., Any]],
 ) -> Service:
     """Imports a function dynamically, and locates the app.
@@ -228,12 +237,9 @@ def import_single_function_service(
     service_deps: Optional[Sequence["modal._object._Object"]] = None
     active_app: modal.app._App
 
-    user_cls_or_cls: typing.Union[None, type, modal.cls.Cls]
-    user_cls_instance = None
-
     if ser_fun is not None:
         # This is a serialized function we already fetched from the server
-        user_cls_or_cls, user_defined_callable = ser_cls, ser_fun
+        user_defined_callable = ser_fun
         active_app = get_active_app_fallback(function_def)
     else:
         # Load the module dynamically
@@ -244,58 +250,22 @@ def import_single_function_service(
             raise LocalFunctionError("Attempted to load a function defined in a function scope")
 
         parts = qual_name.split(".")
-        if len(parts) == 1:
-            # This is a function
-            user_cls_or_cls = None
-            f = getattr(module, qual_name)
-            if isinstance(f, Function):
-                _function: modal._functions._Function[Any, Any, Any] = synchronizer._translate_in(f)  # type: ignore
-                service_deps = _function.deps(only_explicit_mounts=True)
-                user_defined_callable = _function.get_raw_f()
-                assert _function._app  # app should always be set on a decorated function
-                active_app = _function._app
-            else:
-                user_defined_callable = f
-                active_app = get_active_app_fallback(function_def)
-
-        elif len(parts) == 2:
-            # This path should only be triggered by @build class builder methods and can be removed
-            # once @build is deprecated.
-            assert not function_def.use_method_name  # new "placeholder methods" should not be invoked directly!
-            assert function_def.is_builder_function
-            cls_name, fun_name = parts
-            user_cls_or_cls = getattr(module, cls_name)
-            if isinstance(user_cls_or_cls, modal.cls.Cls):
-                # The cls decorator is in global scope
-                _cls = typing.cast(modal.cls._Cls, synchronizer._translate_in(user_cls_or_cls))
-                user_defined_callable = _cls._callables[fun_name]
-                # Intentionally not including these, since @build functions don't actually
-                # forward the information from their parent class.
-                # service_deps = _cls._get_class_service_function().deps(only_explicit_mounts=True)
-                assert _cls._app
-                active_app = _cls._app
-            else:
-                # This is non-decorated class
-                user_defined_callable = getattr(user_cls_or_cls, fun_name)  # unbound method
-                active_app = get_active_app_fallback(function_def)
-        else:
+        if len(parts) != 1:
             raise InvalidError(f"Invalid function qualname {qual_name}")
 
-    # Instantiate the class if it's defined
-    if user_cls_or_cls:
-        if isinstance(user_cls_or_cls, modal.cls.Cls):
-            # This code is only used for @build methods on classes
-            _cls = typing.cast(modal.cls._Cls, user_cls_or_cls)
-            user_cls_instance = get_user_class_instance(_cls, (), {})
-            # Bind the unbound method to the instance as self (using the descriptor protocol!)
+        f = getattr(module, qual_name)
+        if isinstance(f, Function):
+            _function: modal._functions._Function[Any, Any, Any] = synchronizer._translate_in(f)  # type: ignore
+            service_deps = _function.deps(only_explicit_mounts=True)
+            user_defined_callable = _function.get_raw_f()
+            assert _function._app  # app should always be set on a decorated function
+            active_app = _function._app
         else:
-            # serialized=True or "undecorated"
-            user_cls_instance = user_cls_or_cls()
-
-        user_defined_callable = user_defined_callable.__get__(user_cls_instance)
+            # function isn't decorated in global scope
+            user_defined_callable = f
+            active_app = get_active_app_fallback(function_def)
 
     return ImportedFunction(
-        user_cls_instance,
         active_app,
         service_deps,
         user_defined_callable,