modal 0.62.115__py3-none-any.whl → 0.72.13__py3-none-any.whl

modal/gpu.py

@@ -1,10 +1,10 @@
 # Copyright Modal Labs 2022
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from modal_proto import api_pb2
 
-from .exception import InvalidError, deprecation_error, deprecation_warning
+from .exception import InvalidError
 
 
 @dataclass(frozen=True)
@@ -31,7 +31,7 @@ class T4(_GPUConfig):
 
     def __init__(
         self,
-        count: int = 1,  # Number of GPUs per container. Defaults to 1. Useful if you have very large models that don't fit on a single GPU.
+        count: int = 1,  # Number of GPUs per container. Defaults to 1.
     ):
         super().__init__(api_pb2.GPU_TYPE_T4, count, 0)
 
@@ -43,12 +43,13 @@ class L4(_GPUConfig):
     """
     [NVIDIA L4 Tensor Core](https://www.nvidia.com/en-us/data-center/l4/) GPU class.
 
-    A mid-tier data center GPU based on the Ada Lovelace architecture, providing 24GiB of GPU memory. Includes RTX (ray tracing) support.
+    A mid-tier data center GPU based on the Ada Lovelace architecture, providing 24GiB of GPU memory.
+    Includes RTX (ray tracing) support.
     """
 
     def __init__(
         self,
-        count: int = 1,  # Number of GPUs per container. Defaults to 1. Useful if you have very large models that don't fit on a single GPU.
+        count: int = 1,  # Number of GPUs per container. Defaults to 1.
     ):
         super().__init__(api_pb2.GPU_TYPE_L4, count, 0)
 
@@ -66,22 +67,12 @@ class A100(_GPUConfig):
     def __init__(
         self,
         *,
-        count: int = 1,  # Number of GPUs per container. Defaults to 1. Useful if you have very large models that don't fit on a single GPU.
-        memory: int = 0,  # Deprecated. Use `size` instead.
+        count: int = 1,  # Number of GPUs per container. Defaults to 1.
         size: Union[str, None] = None,  # Select GiB configuration of GPU device: "40GB" or "80GB". Defaults to "40GB".
     ):
-        allowed_memory_values = {40, 80}
         allowed_size_values = {"40GB", "80GB"}
-        if memory == 20:
-            raise ValueError(
-                "A100 20GB is unsupported, consider `modal.A10G`, `modal.A100(memory_gb='40')`, or `modal.H100` instead"
-            )
-        elif memory and size:
-            raise ValueError("Cannot specify both `memory` and `size`. Just specify `size`.")
-        elif memory:
-            if memory not in allowed_memory_values:
-                raise ValueError(f"A100s can only have memory values of {allowed_memory_values} => memory={memory}")
-        elif size:
+
+        if size:
             if size not in allowed_size_values:
                 raise ValueError(
                     f"size='{size}' is invalid. A100s can only have memory values of {allowed_size_values}."
@@ -106,14 +97,17 @@ class A10G(_GPUConfig):
     """
     [NVIDIA A10G Tensor Core](https://www.nvidia.com/en-us/data-center/products/a10-gpu/) GPU class.
 
-    A mid-tier data center GPU based on the Ampere architecture, providing 24 GiB of memory. A10G GPUs deliver up to 3.3x better ML training performance, 3x better ML inference performance,
+    A mid-tier data center GPU based on the Ampere architecture, providing 24 GiB of memory.
+    A10G GPUs deliver up to 3.3x better ML training performance, 3x better ML inference performance,
     and 3x better graphics performance, in comparison to NVIDIA T4 GPUs.
     """
 
     def __init__(
         self,
         *,
-        count: int = 1,  # Number of GPUs per container. Defaults to 1. Useful if you have very large models that don't fit on a single GPU.
+        # Number of GPUs per container. Defaults to 1.
+        # Useful if you have very large models that don't fit on a single GPU.
+        count: int = 1,
     ):
         super().__init__(api_pb2.GPU_TYPE_A10G, count)
 
@@ -125,13 +119,17 @@ class H100(_GPUConfig):
     """
     [NVIDIA H100 Tensor Core](https://www.nvidia.com/en-us/data-center/h100/) GPU class.
 
-    The flagship data center GPU of the Hopper architecture. Enhanced support for FP8 precision and a Transformer Engine that provides up to 4X faster training over the prior generation for GPT-3 (175B) models.
+    The flagship data center GPU of the Hopper architecture.
+    Enhanced support for FP8 precision and a Transformer Engine that provides up to 4X faster training
+    over the prior generation for GPT-3 (175B) models.
     """
 
     def __init__(
         self,
         *,
-        count: int = 1,  # Number of GPUs per container. Defaults to 1. Useful if you have very large models that don't fit on a single GPU.
+        # Number of GPUs per container. Defaults to 1.
+        # Useful if you have very large models that don't fit on a single GPU.
+        count: int = 1,
     ):
         super().__init__(api_pb2.GPU_TYPE_H100, count)
 
@@ -139,6 +137,27 @@ class H100(_GPUConfig):
         return f"GPU(H100, count={self.count})"
 
 
+class L40S(_GPUConfig):
+    """
+    [NVIDIA L40S](https://www.nvidia.com/en-us/data-center/l40s/) GPU class.
+
+    The L40S is a data center GPU for the Ada Lovelace architecture. It has 48 GB of on-chip
+    GDDR6 RAM and enhanced support for FP8 precision.
+    """
+
+    def __init__(
+        self,
+        *,
+        # Number of GPUs per container. Defaults to 1.
+        # Useful if you have very large models that don't fit on a single GPU.
+        count: int = 1,
+    ):
+        super().__init__(api_pb2.GPU_TYPE_L40S, count)
+
+    def __repr__(self):
+        return f"GPU(L40S, count={self.count})"
+
+
 class Any(_GPUConfig):
     """Selects any one of the GPU classes available within Modal, according to availability."""
 
@@ -149,24 +168,28 @@ class Any(_GPUConfig):
         return f"GPU(Any, count={self.count})"
 
 
-STRING_TO_GPU_CONFIG = {
+STRING_TO_GPU_CONFIG: dict[str, Callable] = {
     "t4": T4,
     "l4": L4,
     "a100": A100,
+    "a100-80gb": lambda: A100(size="80GB"),
     "h100": H100,
     "a10g": A10G,
+    "l40s": L40S,
     "any": Any,
 }
-display_string_to_config = "\n".join(
-    f'- "{key}" → `{cls()}`' for key, cls in STRING_TO_GPU_CONFIG.items() if key != "inf2"
-)
+display_string_to_config = "\n".join(f'- "{key}" → `{c()}`' for key, c in STRING_TO_GPU_CONFIG.items() if key != "inf2")
 __doc__ = f"""
 **GPU configuration shortcodes**
 
-The following are the valid `str` values for the `gpu` parameter of [`@app.function`](/docs/reference/modal.Stub#function).
+The following are the valid `str` values for the `gpu` parameter of
+[`@app.function`](/docs/reference/modal.App#function).
 
 {display_string_to_config}
 
+The shortcodes also support specifying count by suffixing `:N` to acquire `N` GPUs.
+For example, `a10g:4` will provision 4 A10G GPUs.
+
 Other configurations can be created using the constructors documented below.
 """
 
@@ -174,7 +197,7 @@ Other configurations can be created using the constructors documented below.
 GPU_T = Union[None, bool, str, _GPUConfig]
 
 
-def _parse_gpu_config(value: GPU_T, raise_on_true: bool = True) -> Optional[_GPUConfig]:
+def _parse_gpu_config(value: GPU_T) -> Optional[_GPUConfig]:
     if isinstance(value, _GPUConfig):
         return value
     elif isinstance(value, str):
@@ -187,37 +210,27 @@ def _parse_gpu_config(value: GPU_T, raise_on_true: bool = True) -> Optional[_GPU
                 raise InvalidError(f"Invalid GPU count: {count_str}. Value must be an integer.")
 
         if value.lower() == "a100-20g":
-            return A100(memory=20, count=count)  # Triggers unsupported error underneath.
+            return A100(size="20GB", count=count)  # Triggers unsupported error underneath.
         elif value.lower() == "a100-40gb":
             return A100(size="40GB", count=count)
         elif value.lower() == "a100-80gb":
             return A100(size="80GB", count=count)
         elif value.lower() not in STRING_TO_GPU_CONFIG:
             raise InvalidError(
-                f"Invalid GPU type: {value}. Value must be one of {list(STRING_TO_GPU_CONFIG.keys())} (case-insensitive)."
+                f"Invalid GPU type: {value}. "
+                f"Value must be one of {list(STRING_TO_GPU_CONFIG.keys())} (case-insensitive)."
             )
         else:
             return STRING_TO_GPU_CONFIG[value.lower()](count=count)
-    elif value is True:
-        if raise_on_true:
-            deprecation_error(
-                (2022, 12, 19), 'Setting gpu=True is deprecated. Use `gpu="any"` or `gpu=modal.gpu.Any()` instead.'
-            )
-        else:
-            # We didn't support targeting a GPU type for run_function until 2023-12-12
-            deprecation_warning(
-                (2023, 12, 13), 'Setting gpu=True is deprecated. Use `gpu="any"` or `gpu=modal.gpu.Any()` instead.'
-            )
-        return Any()
     elif value is None or value is False:
         return None
     else:
         raise InvalidError(f"Invalid GPU config: {value}. Value must be a string, a `GPUConfig` object, or `None`.")
 
 
-def parse_gpu_config(value: GPU_T, raise_on_true: bool = True) -> api_pb2.GPUConfig:
+def parse_gpu_config(value: GPU_T) -> api_pb2.GPUConfig:
     """mdmd:hidden"""
-    gpu_config = _parse_gpu_config(value, raise_on_true)
+    gpu_config = _parse_gpu_config(value)
     if gpu_config is None:
         return api_pb2.GPUConfig()
     return gpu_config._to_proto()
@@ -225,4 +238,4 @@ def parse_gpu_config(value: GPU_T, raise_on_true: bool = True) -> api_pb2.GPUCon
 
 def display_gpu_config(value: GPU_T) -> str:
     """mdmd:hidden"""
-    return repr(_parse_gpu_config(value, False))
+    return repr(_parse_gpu_config(value))