modal 0.72.4__py3-none-any.whl → 0.72.48__py3-none-any.whl

modal/functions.py

@@ -26,6 +26,7 @@ from modal_proto import api_pb2
 from modal_proto.modal_api_grpc import ModalClientModal
 
 from ._location import parse_cloud_provider
+from ._object import _get_environment_name, _Object, live_method, live_method_gen
 from ._pty import get_pty_info
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
@@ -71,7 +72,6 @@ from .gpu import GPU_T, parse_gpu_config
 from .image import _Image
 from .mount import _get_client_mount, _Mount, get_auto_mounts
 from .network_file_system import _NetworkFileSystem, network_file_system_mount_protos
-from .object import _get_environment_name, _Object, live_method, live_method_gen
 from .output import _get_output_manager
 from .parallel_map import (
     _for_each_async,
@@ -383,12 +383,15 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     _serve_mounts: frozenset[_Mount]  # set at load time, only by loader
     _app: Optional["modal.app._App"] = None
     _obj: Optional["modal.cls._Obj"] = None  # only set for InstanceServiceFunctions and bound instance methods
-    _web_url: Optional[str]
+
+    _webhook_config: Optional[api_pb2.WebhookConfig] = None  # this is set in definition scope, only locally
+    _web_url: Optional[str]  # this is set on hydration
+
     _function_name: Optional[str]
     _is_method: bool
     _spec: Optional[_FunctionSpec] = None
     _tag: str
-    _raw_f: Callable[..., Any]
+    _raw_f: Optional[Callable[..., Any]]  # this is set to None for a "class service [function]"
     _build_args: dict
 
     _is_generator: Optional[bool] = None
@@ -474,7 +477,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         _experimental_buffer_containers: Optional[int] = None,
         _experimental_proxy_ip: Optional[str] = None,
         _experimental_custom_scaling_factor: Optional[float] = None,
-    ) -> None:
+    ) -> "_Function":
         """mdmd:hidden"""
         # Needed to avoid circular imports
         from .partial_function import _find_partial_methods_for_user_cls, _PartialFunctionFlags
@@ -573,7 +576,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 )
                 image = _Image._from_args(
                     base_images={"base": image},
-                    build_function=snapshot_function,
+                    build_function=snapshot_function,  # type: ignore   # TODO: separate functions.py and _functions.py
                     force_build=image.force_build or pf.force_build,
                 )
 
@@ -785,7 +788,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     task_idle_timeout_secs=container_idle_timeout or 0,
                     concurrency_limit=concurrency_limit or 0,
                     pty_info=pty_info,
-                    cloud_provider=cloud_provider,
+                    cloud_provider=cloud_provider,  # Deprecated at some point
+                    cloud_provider_str=cloud.upper() if cloud else "",  # Supersedes cloud_provider
                     warm_pool_size=keep_warm or 0,
                     runtime=config.get("function_runtime"),
                     runtime_debug=config.get("function_runtime_debug"),
@@ -911,6 +915,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         obj._cluster_size = cluster_size
         obj._is_method = False
         obj._spec = function_spec  # needed for modal shell
+        obj._webhook_config = webhook_config  # only set locally
 
         # Used to check whether we should rebuild a modal.Image which uses `run_function`.
         gpus: list[GPU_T] = gpu if isinstance(gpu, list) else [gpu]
@@ -962,7 +967,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     f"The {identity} has not been hydrated with the metadata it needs to run on Modal{reason}."
                 )
 
-            assert parent._client.stub
+            assert parent._client and parent._client.stub
 
             if can_use_parent:
                 # We can end up here if parent wasn't hydrated when class was instantiated, but has been since.
@@ -983,9 +988,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             else:
                 serialized_params = serialize((args, kwargs))
             environment_name = _get_environment_name(None, resolver)
-            assert parent is not None
+            assert parent is not None and parent.is_hydrated
             req = api_pb2.FunctionBindParamsRequest(
-                function_id=parent._object_id,
+                function_id=parent.object_id,
                 serialized_params=serialized_params,
                 function_options=options,
                 environment_name=environment_name
@@ -1032,11 +1037,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             """
                 )
             )
-        assert self._client and self._client.stub
         request = api_pb2.FunctionUpdateSchedulingParamsRequest(
-            function_id=self._object_id, warm_pool_size_override=warm_pool_size
+            function_id=self.object_id, warm_pool_size_override=warm_pool_size
         )
-        await retry_transient_errors(self._client.stub.FunctionUpdateSchedulingParams, request)
+        await retry_transient_errors(self.client.stub.FunctionUpdateSchedulingParams, request)
 
     @classmethod
     @renamed_parameter((2024, 12, 18), "tag", "name")
@@ -1138,11 +1142,15 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         assert self._spec
         return self._spec
 
+    def _is_web_endpoint(self) -> bool:
+        # only defined in definition scope/locally, and not for class methods at the moment
+        return bool(self._webhook_config and self._webhook_config.type != api_pb2.WEBHOOK_TYPE_UNSPECIFIED)
+
     def get_build_def(self) -> str:
         """mdmd:hidden"""
         # Plaintext source and arg definition for the function, so it's part of the image
         # hash. We can't use the cloudpickle hash because it's not very stable.
-        assert hasattr(self, "_raw_f") and hasattr(self, "_build_args")
+        assert hasattr(self, "_raw_f") and hasattr(self, "_build_args") and self._raw_f is not None
         return f"{inspect.getsource(self._raw_f)}\n{repr(self._build_args)}"
 
     # Live handle methods
@@ -1207,12 +1215,13 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     async def is_generator(self) -> bool:
         """mdmd:hidden"""
         # hacky: kind of like @live_method, but not hydrating if we have the value already from local source
+        # TODO(michael) use a common / lightweight method for handling unhydrated metadata properties
         if self._is_generator is not None:
             # this is set if the function or class is local
             return self._is_generator
 
         # not set - this is a from_name lookup - hydrate
-        await self.resolve()
+        await self.hydrate()
         assert self._is_generator is not None  # should be set now
         return self._is_generator
 
@@ -1248,7 +1257,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             _map_invocation(
                 self,  # type: ignore
                 input_queue,
-                self._client,
+                self.client,
                 order_outputs,
                 return_exceptions,
                 count_update_callback,
@@ -1266,7 +1275,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             self,
             args,
             kwargs,
-            client=self._client,
+            client=self.client,
             function_call_invocation_type=function_call_invocation_type,
         )
 
@@ -1276,7 +1285,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         self, args, kwargs, function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType"
     ) -> _Invocation:
         return await _Invocation.create(
-            self, args, kwargs, client=self._client, function_call_invocation_type=function_call_invocation_type
+            self, args, kwargs, client=self.client, function_call_invocation_type=function_call_invocation_type
         )
 
     @warn_if_generator_is_not_consumed()
@@ -1287,7 +1296,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             self,
             args,
             kwargs,
-            client=self._client,
+            client=self.client,
             function_call_invocation_type=api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC_LEGACY,
         )
         async for res in invocation.run_generator():
@@ -1303,7 +1312,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             self,
             args,
             kwargs,
-            client=self._client,
+            client=self.client,
             function_call_invocation_type=api_pb2.FUNCTION_CALL_INVOCATION_TYPE_ASYNC_LEGACY,
         )
 
@@ -1452,14 +1461,14 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
     def get_raw_f(self) -> Callable[..., Any]:
         """Return the inner Python object wrapped by this Modal Function."""
+        assert self._raw_f is not None
         return self._raw_f
 
     @live_method
     async def get_current_stats(self) -> FunctionStats:
         """Return a `FunctionStats` object describing the current function's queue and runner counts."""
-        assert self._client.stub
         resp = await retry_transient_errors(
-            self._client.stub.FunctionGetCurrentStats,
+            self.client.stub.FunctionGetCurrentStats,
             api_pb2.FunctionGetCurrentStatsRequest(function_id=self.object_id),
             total_timeout=10.0,
         )
@@ -1491,8 +1500,7 @@ class _FunctionCall(typing.Generic[ReturnType], _Object, type_prefix="fc"):
     _is_generator: bool = False
 
     def _invocation(self):
-        assert self._client.stub
-        return _Invocation(self._client.stub, self.object_id, self._client)
+        return _Invocation(self.client.stub, self.object_id, self.client)
 
     async def get(self, timeout: Optional[float] = None) -> ReturnType:
         """Get the result of the function call.

modal/functions.pyi

@@ -1,5 +1,6 @@
 import collections.abc
 import google.protobuf.message
+import modal._object
 import modal._utils.async_utils
 import modal._utils.function_utils
 import modal.app
@@ -133,17 +134,20 @@ ReturnType = typing.TypeVar("ReturnType", covariant=True)
 
 OriginalReturnType = typing.TypeVar("OriginalReturnType", covariant=True)
 
-class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object._Object):
+SUPERSELF = typing.TypeVar("SUPERSELF", covariant=True)
+
+class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal._object._Object):
     _info: typing.Optional[modal._utils.function_utils.FunctionInfo]
     _serve_mounts: frozenset[modal.mount._Mount]
     _app: typing.Optional[modal.app._App]
     _obj: typing.Optional[modal.cls._Obj]
+    _webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig]
     _web_url: typing.Optional[str]
     _function_name: typing.Optional[str]
     _is_method: bool
     _spec: typing.Optional[_FunctionSpec]
     _tag: str
-    _raw_f: typing.Callable[..., typing.Any]
+    _raw_f: typing.Optional[collections.abc.Callable[..., typing.Any]]
     _build_args: dict
     _is_generator: typing.Optional[bool]
     _cluster_size: typing.Optional[int]
@@ -197,7 +201,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
         _experimental_buffer_containers: typing.Optional[int] = None,
         _experimental_proxy_ip: typing.Optional[str] = None,
         _experimental_custom_scaling_factor: typing.Optional[float] = None,
-    ) -> None: ...
+    ) -> _Function: ...
     def _bind_parameters(
         self,
         obj: modal.cls._Obj,
@@ -228,6 +232,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
     def info(self) -> modal._utils.function_utils.FunctionInfo: ...
     @property
     def spec(self) -> _FunctionSpec: ...
+    def _is_web_endpoint(self) -> bool: ...
     def get_build_def(self) -> str: ...
     def _initialize_from_empty(self): ...
     def _hydrate_metadata(self, metadata: typing.Optional[google.protobuf.message.Message]): ...
@@ -254,10 +259,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
     def local(self, *args: P.args, **kwargs: P.kwargs) -> OriginalReturnType: ...
     async def _experimental_spawn(self, *args: P.args, **kwargs: P.kwargs) -> _FunctionCall[ReturnType]: ...
     async def spawn(self, *args: P.args, **kwargs: P.kwargs) -> _FunctionCall[ReturnType]: ...
-    def get_raw_f(self) -> typing.Callable[..., typing.Any]: ...
+    def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]: ...
     async def get_current_stats(self) -> FunctionStats: ...
 
-    class __map_spec(typing_extensions.Protocol):
+    class __map_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, *input_iterators, kwargs={}, order_outputs: bool = True, return_exceptions: bool = False
         ) -> modal._utils.async_utils.AsyncOrSyncIterable: ...
@@ -269,9 +274,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
             return_exceptions: bool = False,
         ) -> typing.AsyncGenerator[typing.Any, None]: ...
 
-    map: __map_spec
+    map: __map_spec[typing_extensions.Self]
 
-    class __starmap_spec(typing_extensions.Protocol):
+    class __starmap_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             input_iterator: typing.Iterable[typing.Sequence[typing.Any]],
@@ -289,13 +294,13 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
             return_exceptions: bool = False,
         ) -> typing.AsyncIterable[typing.Any]: ...
 
-    starmap: __starmap_spec
+    starmap: __starmap_spec[typing_extensions.Self]
 
-    class __for_each_spec(typing_extensions.Protocol):
+    class __for_each_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
         async def aio(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
 
-    for_each: __for_each_spec
+    for_each: __for_each_spec[typing_extensions.Self]
 
 ReturnType_INNER = typing.TypeVar("ReturnType_INNER", covariant=True)
 
@@ -306,12 +311,13 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     _serve_mounts: frozenset[modal.mount.Mount]
     _app: typing.Optional[modal.app.App]
     _obj: typing.Optional[modal.cls.Obj]
+    _webhook_config: typing.Optional[modal_proto.api_pb2.WebhookConfig]
     _web_url: typing.Optional[str]
     _function_name: typing.Optional[str]
     _is_method: bool
     _spec: typing.Optional[_FunctionSpec]
     _tag: str
-    _raw_f: typing.Callable[..., typing.Any]
+    _raw_f: typing.Optional[collections.abc.Callable[..., typing.Any]]
     _build_args: dict
     _is_generator: typing.Optional[bool]
     _cluster_size: typing.Optional[int]
@@ -366,7 +372,7 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
         _experimental_buffer_containers: typing.Optional[int] = None,
         _experimental_proxy_ip: typing.Optional[str] = None,
         _experimental_custom_scaling_factor: typing.Optional[float] = None,
-    ) -> None: ...
+    ) -> Function: ...
     def _bind_parameters(
         self,
         obj: modal.cls.Obj,
@@ -375,11 +381,11 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
         kwargs: dict[str, typing.Any],
     ) -> Function: ...
 
-    class __keep_warm_spec(typing_extensions.Protocol):
+    class __keep_warm_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, warm_pool_size: int) -> None: ...
         async def aio(self, warm_pool_size: int) -> None: ...
 
-    keep_warm: __keep_warm_spec
+    keep_warm: __keep_warm_spec[typing_extensions.Self]
 
     @classmethod
     def from_name(
@@ -416,6 +422,7 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     def info(self) -> modal._utils.function_utils.FunctionInfo: ...
     @property
     def spec(self) -> _FunctionSpec: ...
+    def _is_web_endpoint(self) -> bool: ...
     def get_build_def(self) -> str: ...
     def _initialize_from_empty(self): ...
     def _hydrate_metadata(self, metadata: typing.Optional[google.protobuf.message.Message]): ...
@@ -428,7 +435,7 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     @property
     def cluster_size(self) -> int: ...
 
-    class ___map_spec(typing_extensions.Protocol):
+    class ___map_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, input_queue: modal.parallel_map.SynchronizedQueue, order_outputs: bool, return_exceptions: bool
         ) -> typing.Generator[typing.Any, None, None]: ...
@@ -436,70 +443,70 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
             self, input_queue: modal.parallel_map.SynchronizedQueue, order_outputs: bool, return_exceptions: bool
         ) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
-    _map: ___map_spec
+    _map: ___map_spec[typing_extensions.Self]
 
-    class ___call_function_spec(typing_extensions.Protocol[ReturnType_INNER]):
+    class ___call_function_spec(typing_extensions.Protocol[ReturnType_INNER, SUPERSELF]):
         def __call__(self, args, kwargs) -> ReturnType_INNER: ...
         async def aio(self, args, kwargs) -> ReturnType_INNER: ...
 
-    _call_function: ___call_function_spec[ReturnType]
+    _call_function: ___call_function_spec[ReturnType, typing_extensions.Self]
 
-    class ___call_function_nowait_spec(typing_extensions.Protocol):
+    class ___call_function_nowait_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, args, kwargs, function_call_invocation_type: int) -> _Invocation: ...
         async def aio(self, args, kwargs, function_call_invocation_type: int) -> _Invocation: ...
 
-    _call_function_nowait: ___call_function_nowait_spec
+    _call_function_nowait: ___call_function_nowait_spec[typing_extensions.Self]
 
-    class ___call_generator_spec(typing_extensions.Protocol):
+    class ___call_generator_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, args, kwargs): ...
         def aio(self, args, kwargs): ...
 
-    _call_generator: ___call_generator_spec
+    _call_generator: ___call_generator_spec[typing_extensions.Self]
 
-    class ___call_generator_nowait_spec(typing_extensions.Protocol):
+    class ___call_generator_nowait_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, args, kwargs): ...
         async def aio(self, args, kwargs): ...
 
-    _call_generator_nowait: ___call_generator_nowait_spec
+    _call_generator_nowait: ___call_generator_nowait_spec[typing_extensions.Self]
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_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[P, ReturnType]
+    remote: __remote_spec[ReturnType, P, typing_extensions.Self]
 
-    class __remote_gen_spec(typing_extensions.Protocol):
+    class __remote_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
         def aio(self, *args, **kwargs) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
-    remote_gen: __remote_gen_spec
+    remote_gen: __remote_gen_spec[typing_extensions.Self]
 
     def _is_local(self): ...
     def _get_info(self) -> modal._utils.function_utils.FunctionInfo: ...
     def _get_obj(self) -> typing.Optional[modal.cls.Obj]: ...
     def local(self, *args: P.args, **kwargs: P.kwargs) -> OriginalReturnType: ...
 
-    class ___experimental_spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    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]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    _experimental_spawn: ___experimental_spawn_spec[P, ReturnType]
+    _experimental_spawn: ___experimental_spawn_spec[ReturnType, P, typing_extensions.Self]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    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]: ...
         async def aio(self, *args: P_INNER.args, **kwargs: P_INNER.kwargs) -> FunctionCall[ReturnType_INNER]: ...
 
-    spawn: __spawn_spec[P, ReturnType]
+    spawn: __spawn_spec[ReturnType, P, typing_extensions.Self]
 
-    def get_raw_f(self) -> typing.Callable[..., typing.Any]: ...
+    def get_raw_f(self) -> collections.abc.Callable[..., typing.Any]: ...
 
-    class __get_current_stats_spec(typing_extensions.Protocol):
+    class __get_current_stats_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self) -> FunctionStats: ...
         async def aio(self) -> FunctionStats: ...
 
-    get_current_stats: __get_current_stats_spec
+    get_current_stats: __get_current_stats_spec[typing_extensions.Self]
 
-    class __map_spec(typing_extensions.Protocol):
+    class __map_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, *input_iterators, kwargs={}, order_outputs: bool = True, return_exceptions: bool = False
         ) -> modal._utils.async_utils.AsyncOrSyncIterable: ...
@@ -511,9 +518,9 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
             return_exceptions: bool = False,
         ) -> typing.AsyncGenerator[typing.Any, None]: ...
 
-    map: __map_spec
+    map: __map_spec[typing_extensions.Self]
 
-    class __starmap_spec(typing_extensions.Protocol):
+    class __starmap_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self,
             input_iterator: typing.Iterable[typing.Sequence[typing.Any]],
@@ -531,15 +538,15 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
             return_exceptions: bool = False,
         ) -> typing.AsyncIterable[typing.Any]: ...
 
-    starmap: __starmap_spec
+    starmap: __starmap_spec[typing_extensions.Self]
 
-    class __for_each_spec(typing_extensions.Protocol):
+    class __for_each_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
         async def aio(self, *input_iterators, kwargs={}, ignore_exceptions: bool = False): ...
 
-    for_each: __for_each_spec
+    for_each: __for_each_spec[typing_extensions.Self]
 
-class _FunctionCall(typing.Generic[ReturnType], modal.object._Object):
+class _FunctionCall(typing.Generic[ReturnType], modal._object._Object):
     _is_generator: bool
 
     def _invocation(self): ...
@@ -558,29 +565,29 @@ class FunctionCall(typing.Generic[ReturnType], modal.object.Object):
     def __init__(self, *args, **kwargs): ...
     def _invocation(self): ...
 
-    class __get_spec(typing_extensions.Protocol[ReturnType_INNER]):
+    class __get_spec(typing_extensions.Protocol[ReturnType_INNER, SUPERSELF]):
         def __call__(self, timeout: typing.Optional[float] = None) -> ReturnType_INNER: ...
         async def aio(self, timeout: typing.Optional[float] = None) -> ReturnType_INNER: ...
 
-    get: __get_spec[ReturnType]
+    get: __get_spec[ReturnType, typing_extensions.Self]
 
-    class __get_gen_spec(typing_extensions.Protocol):
+    class __get_gen_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self) -> typing.Generator[typing.Any, None, None]: ...
         def aio(self) -> collections.abc.AsyncGenerator[typing.Any, None]: ...
 
-    get_gen: __get_gen_spec
+    get_gen: __get_gen_spec[typing_extensions.Self]
 
-    class __get_call_graph_spec(typing_extensions.Protocol):
+    class __get_call_graph_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self) -> list[modal.call_graph.InputInfo]: ...
         async def aio(self) -> list[modal.call_graph.InputInfo]: ...
 
-    get_call_graph: __get_call_graph_spec
+    get_call_graph: __get_call_graph_spec[typing_extensions.Self]
 
-    class __cancel_spec(typing_extensions.Protocol):
+    class __cancel_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(self, terminate_containers: bool = False): ...
         async def aio(self, terminate_containers: bool = False): ...
 
-    cancel: __cancel_spec
+    cancel: __cancel_spec[typing_extensions.Self]
 
     class __from_id_spec(typing_extensions.Protocol):
         def __call__(

modal/gpu.py

@@ -9,8 +9,9 @@ from .exception import InvalidError
 
 @dataclass(frozen=True)
 class _GPUConfig:
-    type: "api_pb2.GPUType.V"
+    type: "api_pb2.GPUType.V"  # Deprecated, at some point
     count: int
+    gpu_type: str
     memory: int = 0
 
     def _to_proto(self) -> api_pb2.GPUConfig:
@@ -19,6 +20,7 @@ class _GPUConfig:
             type=self.type,
             count=self.count,
             memory=self.memory,
+            gpu_type=self.gpu_type,
         )
 
 
@@ -26,14 +28,14 @@ class T4(_GPUConfig):
     """
     [NVIDIA T4 Tensor Core](https://www.nvidia.com/en-us/data-center/tesla-t4/) GPU class.
 
-    A low-cost data center GPU based on the Turing architecture, providing 16GiB of GPU memory.
+    A low-cost data center GPU based on the Turing architecture, providing 16GB of GPU memory.
     """
 
     def __init__(
         self,
         count: int = 1,  # Number of GPUs per container. Defaults to 1.
     ):
-        super().__init__(api_pb2.GPU_TYPE_T4, count, 0)
+        super().__init__(api_pb2.GPU_TYPE_T4, count, "T4")
 
     def __repr__(self):
         return f"GPU(T4, count={self.count})"
@@ -43,7 +45,7 @@ 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.
+    A mid-tier data center GPU based on the Ada Lovelace architecture, providing 24GB of GPU memory.
     Includes RTX (ray tracing) support.
     """
 
@@ -51,7 +53,7 @@ class L4(_GPUConfig):
         self,
         count: int = 1,  # Number of GPUs per container. Defaults to 1.
     ):
-        super().__init__(api_pb2.GPU_TYPE_L4, count, 0)
+        super().__init__(api_pb2.GPU_TYPE_L4, count, "L4")
 
     def __repr__(self):
         return f"GPU(L4, count={self.count})"
@@ -61,30 +63,21 @@ class A100(_GPUConfig):
     """
     [NVIDIA A100 Tensor Core](https://www.nvidia.com/en-us/data-center/a100/) GPU class.
 
-    The flagship data center GPU of the Ampere architecture. Available in 40GiB and 80GiB GPU memory configurations.
+    The flagship data center GPU of the Ampere architecture. Available in 40GB and 80GB GPU memory configurations.
     """
 
     def __init__(
         self,
         *,
         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".
+        size: Union[str, None] = None,  # Select GB configuration of GPU device: "40GB" or "80GB". Defaults to "40GB".
     ):
-        allowed_size_values = {"40GB", "80GB"}
-
-        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}."
-                )
-            memory = int(size.replace("GB", ""))
+        if size == "40GB" or not size:
+            super().__init__(api_pb2.GPU_TYPE_A100, count, "A100-40GB", 40)
+        elif size == "80GB":
+            super().__init__(api_pb2.GPU_TYPE_A100_80GB, count, "A100-80GB", 80)
         else:
-            memory = 40
-
-        if memory == 80:
-            super().__init__(api_pb2.GPU_TYPE_A100_80GB, count, memory)
-        else:
-            super().__init__(api_pb2.GPU_TYPE_A100, count, memory)
+            raise ValueError(f"size='{size}' is invalid. A100s can only have memory values of 40GB or 80GB.")
 
     def __repr__(self):
         if self.memory == 80:
@@ -97,7 +90,7 @@ 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.
+    A mid-tier data center GPU based on the Ampere architecture, providing 24 GB 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.
     """
@@ -109,7 +102,7 @@ class A10G(_GPUConfig):
         # 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)
+        super().__init__(api_pb2.GPU_TYPE_A10G, count, "A10G")
 
     def __repr__(self):
         return f"GPU(A10G, count={self.count})"
@@ -131,7 +124,7 @@ class H100(_GPUConfig):
         # 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)
+        super().__init__(api_pb2.GPU_TYPE_H100, count, "H100")
 
     def __repr__(self):
         return f"GPU(H100, count={self.count})"
@@ -152,7 +145,7 @@ class L40S(_GPUConfig):
         # 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)
+        super().__init__(api_pb2.GPU_TYPE_L40S, count, "L40S")
 
     def __repr__(self):
         return f"GPU(L40S, count={self.count})"
@@ -162,7 +155,7 @@ class Any(_GPUConfig):
     """Selects any one of the GPU classes available within Modal, according to availability."""
 
     def __init__(self, *, count: int = 1):
-        super().__init__(api_pb2.GPU_TYPE_ANY, count)
+        super().__init__(api_pb2.GPU_TYPE_ANY, count, "ANY")
 
     def __repr__(self):
         return f"GPU(Any, count={self.count})"