modal 1.1.5.dev50__tar.gz → 1.1.5.dev52__tar.gz

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/PKG-INFO

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

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal/_runtime/gpu_memory_snapshot.py

@@ -18,6 +18,12 @@ 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 [1].
@@ -44,7 +50,7 @@ class CudaCheckpointProcess:
     pid: int
     state: CudaCheckpointState
 
-    def toggle(self, target_state: CudaCheckpointState, timeout_secs: float = 5 * 60.0) -> None:
+    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.
         """
@@ -54,7 +60,11 @@ class CudaCheckpointProcess:
         retry_count = 0
         max_retries = 3
 
-        while self._should_continue_toggle(target_state, start_time, timeout_secs):
+        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
@@ -73,10 +83,11 @@ class CudaCheckpointProcess:
         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
+        self, target_state: CudaCheckpointState, start_time: float, refresh: bool = True
     ) -> bool:
         """Check if toggle operation should continue based on current state and timeout."""
-        self.refresh_state()
+        if refresh:
+            self.refresh_state()
 
         if self.state == target_state:
             return False
@@ -85,7 +96,7 @@ class CudaCheckpointProcess:
             raise CudaCheckpointException(f"PID: {self.pid} CUDA process state is {self.state}")
 
         elapsed = time.monotonic() - start_time
-        if elapsed >= timeout_secs:
+        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}"
@@ -101,7 +112,7 @@ class CudaCheckpointProcess:
                 check=True,
                 capture_output=True,
                 text=True,
-                timeout=30,
+                timeout=CUDA_CHECKPOINT_TIMEOUT,
             )
             logger.debug(f"PID: {self.pid} Successfully toggled CUDA checkpoint state")
         except subprocess.CalledProcessError as e:
@@ -121,7 +132,7 @@ class CudaCheckpointProcess:
                 check=True,
                 capture_output=True,
                 text=True,
-                timeout=10,
+                timeout=CUDA_CHECKPOINT_TIMEOUT,
             )
 
             state_str = result.stdout.strip().lower()
@@ -190,6 +201,7 @@ class CudaCheckpointSession:
                 [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,
             )
 
@@ -256,20 +268,11 @@ class CudaCheckpointSession:
             logger.debug("No CUDA sessions to restore.")
             return
 
-        # Validate all states first
-        for proc in self.cuda_processes:
-            proc.refresh_state()  # Refresh state before validation
-            if proc.state != CudaCheckpointState.CHECKPOINTED:
-                raise CudaCheckpointException(
-                    f"PID {proc.pid}: CUDA session not in {CudaCheckpointState.CHECKPOINTED.value} state. "
-                    f"Current state: {proc.state.value}"
-                )
-
         # See checkpoint() for rationale about parallelism.
         start = time.perf_counter()
 
         def restore_process(proc: CudaCheckpointProcess) -> None:
-            proc.toggle(CudaCheckpointState.RUNNING)
+            proc.toggle(CudaCheckpointState.RUNNING, skip_first_refresh=True)
 
         with ThreadPoolExecutor() as executor:
             futures = [executor.submit(restore_process, proc) for proc in self.cuda_processes]

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal/client.pyi

@@ -33,7 +33,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.1.5.dev50",
+        version: str = "1.1.5.dev52",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.
@@ -164,7 +164,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[tuple[str, str]],
-        version: str = "1.1.5.dev50",
+        version: str = "1.1.5.dev52",
     ):
         """mdmd:hidden
         The Modal client object is not intended to be instantiated directly by users.

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal/functions.pyi

@@ -450,7 +450,7 @@ class Function(
 
     _call_generator: ___call_generator_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> ReturnType_INNER:
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
@@ -459,7 +459,7 @@ class Function(
             """Calls the function remotely, executing it with the given arguments and returning the execution's result."""
             ...
 
-    remote: __remote_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    remote: __remote_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, /, *args, **kwargs) -> typing.Generator[typing.Any, None, None]:
@@ -486,7 +486,7 @@ class Function(
         """
         ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class ___experimental_spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """[Experimental] Calls the function with the given arguments, without waiting for the results.
 
@@ -510,7 +510,7 @@ class Function(
             ...
 
     _experimental_spawn: ___experimental_spawn_spec[
-        modal._functions.P, modal._functions.ReturnType, typing_extensions.Self
+        modal._functions.ReturnType, modal._functions.P, typing_extensions.Self
     ]
 
     class ___spawn_map_inner_spec(typing_extensions.Protocol[P_INNER, SUPERSELF]):
@@ -519,7 +519,7 @@ class Function(
 
     _spawn_map_inner: ___spawn_map_inner_spec[modal._functions.P, typing_extensions.Self]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER, SUPERSELF]):
+    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER, SUPERSELF]):
         def __call__(self, /, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]:
             """Calls the function with the given arguments, without waiting for the results.
 
@@ -540,7 +540,7 @@ class Function(
             """
             ...
 
-    spawn: __spawn_spec[modal._functions.P, modal._functions.ReturnType, typing_extensions.Self]
+    spawn: __spawn_spec[modal._functions.ReturnType, modal._functions.P, typing_extensions.Self]
 
     def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]:
         """Return the inner Python object wrapped by this Modal Function."""

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal.egg-info/PKG-INFO

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

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal_proto/api.proto

@@ -1345,6 +1345,14 @@ message FlashProxyUpstreamRequest {
   double timestamp = 2; 
 }
 
+message FlashSetTargetSlotsMetricsRequest {
+  // TODO(claudia): add other metrics to use in autoscaling decisions
+  string function_id = 1;
+  uint32 target_slots = 2;
+}
+
+message FlashSetTargetSlotsMetricsResponse {}
+
 message Function {
   string module_name = 1;
   string function_name = 2;
@@ -3629,6 +3637,7 @@ service ModalClient {
   rpc FlashContainerDeregister(FlashContainerDeregisterRequest) returns (google.protobuf.Empty);
   rpc FlashContainerList(FlashContainerListRequest) returns (FlashContainerListResponse);
   rpc FlashContainerRegister(FlashContainerRegisterRequest) returns (FlashContainerRegisterResponse);
+  rpc FlashSetTargetSlotsMetrics(FlashSetTargetSlotsMetricsRequest) returns (FlashSetTargetSlotsMetricsResponse);
 
   // Functions
   rpc FunctionAsyncInvoke(FunctionAsyncInvokeRequest) returns (FunctionAsyncInvokeResponse);

{modal-1.1.5.dev50 → modal-1.1.5.dev52}/modal_proto/api_grpc.py

@@ -262,6 +262,10 @@ class ModalClientBase(abc.ABC):
     async def FlashContainerRegister(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.FlashContainerRegisterRequest, modal_proto.api_pb2.FlashContainerRegisterResponse]') -> None:
         pass
 
+    @abc.abstractmethod
+    async def FlashSetTargetSlotsMetrics(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.FlashSetTargetSlotsMetricsRequest, modal_proto.api_pb2.FlashSetTargetSlotsMetricsResponse]') -> None:
+        pass
+
     @abc.abstractmethod
     async def FunctionAsyncInvoke(self, stream: 'grpclib.server.Stream[modal_proto.api_pb2.FunctionAsyncInvokeRequest, modal_proto.api_pb2.FunctionAsyncInvokeResponse]') -> None:
         pass
@@ -1070,6 +1074,12 @@ class ModalClientBase(abc.ABC):
                 modal_proto.api_pb2.FlashContainerRegisterRequest,
                 modal_proto.api_pb2.FlashContainerRegisterResponse,
             ),
+            '/modal.client.ModalClient/FlashSetTargetSlotsMetrics': grpclib.const.Handler(
+                self.FlashSetTargetSlotsMetrics,
+                grpclib.const.Cardinality.UNARY_UNARY,
+                modal_proto.api_pb2.FlashSetTargetSlotsMetricsRequest,
+                modal_proto.api_pb2.FlashSetTargetSlotsMetricsResponse,
+            ),
             '/modal.client.ModalClient/FunctionAsyncInvoke': grpclib.const.Handler(
                 self.FunctionAsyncInvoke,
                 grpclib.const.Cardinality.UNARY_UNARY,
@@ -2102,6 +2112,12 @@ class ModalClientStub:
             modal_proto.api_pb2.FlashContainerRegisterRequest,
             modal_proto.api_pb2.FlashContainerRegisterResponse,
         )
+        self.FlashSetTargetSlotsMetrics = grpclib.client.UnaryUnaryMethod(
+            channel,
+            '/modal.client.ModalClient/FlashSetTargetSlotsMetrics',
+            modal_proto.api_pb2.FlashSetTargetSlotsMetricsRequest,
+            modal_proto.api_pb2.FlashSetTargetSlotsMetricsResponse,
+        )
         self.FunctionAsyncInvoke = grpclib.client.UnaryUnaryMethod(
             channel,
             '/modal.client.ModalClient/FunctionAsyncInvoke',