modal 0.66.17__py3-none-any.whl → 0.66.44__py3-none-any.whl

modal/cli/import_refs.py

@@ -19,7 +19,7 @@ from rich.console import Console
 from rich.markdown import Markdown
 
 from modal.app import App, LocalEntrypoint
-from modal.exception import InvalidError, _CliUserExecutionError, deprecation_warning
+from modal.exception import InvalidError, _CliUserExecutionError
 from modal.functions import Function
 
 
@@ -79,7 +79,7 @@ def import_file_or_module(file_or_module: str):
     return module
 
 
-def get_by_object_path(obj: Any, obj_path: Optional[str]) -> Optional[Any]:
+def get_by_object_path(obj: Any, obj_path: str) -> Optional[Any]:
     # Try to evaluate a `.`-delimited object path in a Modal context
     # With the caveat that some object names can actually have `.` in their name (lifecycled methods' tags)
 
@@ -107,35 +107,6 @@ def get_by_object_path(obj: Any, obj_path: Optional[str]) -> Optional[Any]:
     return obj
 
 
-def get_by_object_path_try_possible_app_names(obj: Any, obj_path: Optional[str]) -> Optional[Any]:
-    """This just exists as a dumb workaround to support both "stub" and "app" """
-
-    if obj_path:
-        return get_by_object_path(obj, obj_path)
-    else:
-        app = get_by_object_path(obj, DEFAULT_APP_NAME)
-        stub = get_by_object_path(obj, "stub")
-        if isinstance(app, App):
-            return app
-        elif app is not None and isinstance(stub, App):
-            deprecation_warning(
-                (2024, 4, 20),
-                "The symbol `app` is present at the module level but it's not a Modal app."
-                " We will use `stub` instead, but this will not work in future Modal versions."
-                " Suggestion: change the name of `app` to something else.",
-            )
-            return stub
-        elif isinstance(stub, App):
-            deprecation_warning(
-                (2024, 5, 1),
-                "The symbol `app` is not present but `stub` is. This will not work in future"
-                " Modal versions. Suggestion: change the name of `stub` to `app`.",
-            )
-            return stub
-        else:
-            return None
-
-
 def _infer_function_or_help(
     app: App, module, accept_local_entrypoint: bool, accept_webhook: bool
 ) -> Union[Function, LocalEntrypoint]:
@@ -210,7 +181,7 @@ def import_app(app_ref: str) -> App:
     import_ref = parse_import_ref(app_ref)
 
     module = import_file_or_module(import_ref.file_or_module)
-    app = get_by_object_path_try_possible_app_names(module, import_ref.object_path)
+    app = get_by_object_path(module, import_ref.object_path or DEFAULT_APP_NAME)
 
     if app is None:
         _show_no_auto_detectable_app(import_ref)
@@ -258,7 +229,7 @@ def import_function(
     import_ref = parse_import_ref(func_ref)
 
     module = import_file_or_module(import_ref.file_or_module)
-    app_or_function = get_by_object_path_try_possible_app_names(module, import_ref.object_path)
+    app_or_function = get_by_object_path(module, import_ref.object_path or DEFAULT_APP_NAME)
 
     if app_or_function is None:
         _show_function_ref_help(import_ref, base_cmd)
@@ -279,8 +250,3 @@ def import_function(
         return app_or_function
     else:
         raise click.UsageError(f"{app_or_function} is not a Modal entity (should be an App or Function)")
-
-
-# For backwards compatibility - delete soon
-# We use it in our internal intergration tests
-import_stub = import_app

modal/client.pyi

@@ -31,7 +31,7 @@ class _Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[typing.Tuple[str, str]],
-        version: str = "0.66.17",
+        version: str = "0.66.44",
     ): ...
     def is_closed(self) -> bool: ...
     @property
@@ -90,7 +90,7 @@ class Client:
         server_url: str,
         client_type: int,
         credentials: typing.Optional[typing.Tuple[str, str]],
-        version: str = "0.66.17",
+        version: str = "0.66.44",
     ): ...
     def is_closed(self) -> bool: ...
     @property

modal/cls.py

@@ -113,7 +113,7 @@ class _Obj:
                 method = self._instance_service_function._bind_instance_method(class_bound_method)
                 self._method_functions[method_name] = method
         else:
-            # <v0.63 classes - bind each individual method to the new parameters
+            # looked up <v0.63 classes - bind each individual method to the new parameters
             self._instance_service_function = None
             for method_name, class_bound_method in classbound_methods.items():
                 method = class_bound_method._bind_parameters(self, from_other_workspace, options, args, kwargs)
@@ -125,12 +125,14 @@ class _Obj:
         self._user_cls = user_cls
         self._construction_args = (args, kwargs)  # used for lazy construction in case of explicit constructors
 
-    def _user_cls_instance_constr(self):
+    def _new_user_cls_instance(self):
         args, kwargs = self._construction_args
         if not _use_annotation_parameters(self._user_cls):
             # TODO(elias): deprecate this code path eventually
             user_cls_instance = self._user_cls(*args, **kwargs)
         else:
+            # ignore constructor (assumes there is no custom constructor,
+            # which is guaranteed by _use_annotation_parameters)
             # set the attributes on the class corresponding to annotations
             # with = parameter() specifications
             sig = _get_class_constructor_signature(self._user_cls)
@@ -139,6 +141,7 @@ class _Obj:
             user_cls_instance = self._user_cls.__new__(self._user_cls)  # new instance without running __init__
             user_cls_instance.__dict__.update(bound_vars.arguments)
 
+        # TODO: always use Obj instances instead of making modifications to user cls
         user_cls_instance._modal_functions = self._method_functions  # Needed for PartialFunction.__get__
         return user_cls_instance
 
@@ -163,10 +166,12 @@ class _Obj:
             )
         await self._instance_service_function.keep_warm(warm_pool_size)
 
-    def _get_user_cls_instance(self):
-        """Construct local object lazily. Used for .local() calls."""
+    def _cached_user_cls_instance(self):
+        """Get or construct the local object
+
+        Used for .local() calls and getting attributes of classes"""
         if not self._user_cls_instance:
-            self._user_cls_instance = self._user_cls_instance_constr()  # Instantiate object
+            self._user_cls_instance = self._new_user_cls_instance()  # Instantiate object
 
         return self._user_cls_instance
 
@@ -196,7 +201,7 @@ class _Obj:
     @synchronizer.nowrap
     async def aenter(self):
         if not self.entered:
-            user_cls_instance = self._get_user_cls_instance()
+            user_cls_instance = self._cached_user_cls_instance()
             if hasattr(user_cls_instance, "__aenter__"):
                 await user_cls_instance.__aenter__()
             elif hasattr(user_cls_instance, "__enter__"):
@@ -205,20 +210,22 @@ class _Obj:
 
     def __getattr__(self, k):
         if k in self._method_functions:
-            # if we know the user is accessing a method, we don't have to create an instance
-            # yet, since the user might just call `.remote()` on it which doesn't require
-            # a local instance (in case __init__ does stuff that can't locally)
+            # If we know the user is accessing a *method* and not another attribute,
+            # we don't have to create an instance of the user class yet.
+            # This is because it might just be a call to `.remote()` on it which
+            # doesn't require a local instance.
+            # As long as we have the service function or params, we can do remote calls
+            # without calling the constructor of the class in the calling context.
             return self._method_functions[k]
-        elif self._user_cls_instance_constr:
-            # if it's *not* a method
-            # TODO: To get lazy loading (from_name) of classes to work, we need to avoid
-            #  this path, otherwise local initialization will happen regardless if user
-            #  only runs .remote(), since we don't know methods for the class until we
-            #  load it
-            user_cls_instance = self._get_user_cls_instance()
-            return getattr(user_cls_instance, k)
-        else:
-            raise AttributeError(k)
+
+        # if it's *not* a method, it *might* be an attribute of the class,
+        # so we construct it and proxy the attribute
+        # TODO: To get lazy loading (from_name) of classes to work, we need to avoid
+        #  this path, otherwise local initialization will happen regardless if user
+        #  only runs .remote(), since we don't know methods for the class until we
+        #  load it
+        user_cls_instance = self._cached_user_cls_instance()
+        return getattr(user_cls_instance, k)
 
 
 Obj = synchronize_api(_Obj)

modal/cls.pyi

@@ -37,9 +37,9 @@ class _Obj:
         args,
         kwargs,
     ): ...
-    def _user_cls_instance_constr(self): ...
+    def _new_user_cls_instance(self): ...
     async def keep_warm(self, warm_pool_size: int) -> None: ...
-    def _get_user_cls_instance(self): ...
+    def _cached_user_cls_instance(self): ...
     def enter(self): ...
     @property
     def entered(self): ...
@@ -66,7 +66,7 @@ class Obj:
         kwargs,
     ): ...
     def _uses_common_service_function(self): ...
-    def _user_cls_instance_constr(self): ...
+    def _new_user_cls_instance(self): ...
 
     class __keep_warm_spec(typing_extensions.Protocol):
         def __call__(self, warm_pool_size: int) -> None: ...
@@ -74,7 +74,7 @@ class Obj:
 
     keep_warm: __keep_warm_spec
 
-    def _get_user_cls_instance(self): ...
+    def _cached_user_cls_instance(self): ...
     def enter(self): ...
     @property
     def entered(self): ...

modal/dict.py

@@ -143,12 +143,6 @@ class _Dict(_Object, type_prefix="di"):
 
         return _Dict._from_loader(_load, "Dict()", is_another_app=True, hydrate_lazily=True)
 
-    @staticmethod
-    def persisted(label: str, namespace=api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE, environment_name: Optional[str] = None):
-        """mdmd:hidden"""
-        message = "`Dict.persisted` is deprecated. Please use `Dict.from_name(name, create_if_missing=True)` instead."
-        deprecation_error((2024, 3, 1), message)
-
     @staticmethod
     async def lookup(
         label: str,

modal/dict.pyi

@@ -27,8 +27,6 @@ class _Dict(modal.object._Object):
         create_if_missing: bool = False,
     ) -> _Dict: ...
     @staticmethod
-    def persisted(label: str, namespace=1, environment_name: typing.Optional[str] = None): ...
-    @staticmethod
     async def lookup(
         label: str,
         data: typing.Optional[dict] = None,
@@ -79,8 +77,6 @@ class Dict(modal.object.Object):
         environment_name: typing.Optional[str] = None,
         create_if_missing: bool = False,
     ) -> Dict: ...
-    @staticmethod
-    def persisted(label: str, namespace=1, environment_name: typing.Optional[str] = None): ...
 
     class __lookup_spec(typing_extensions.Protocol):
         def __call__(

modal/experimental.py

@@ -17,7 +17,6 @@ from .partial_function import _PartialFunction, _PartialFunctionFlags
 def stop_fetching_inputs():
     """Don't fetch any more inputs from the server, after the current one.
     The container will exit gracefully after the current input is processed."""
-
     _ContainerIOManager.stop_fetching_inputs()
 
 
@@ -25,7 +24,6 @@ def get_local_input_concurrency():
     """Get the container's local input concurrency.
     If recently reduced to particular value, it can return a larger number than
     set due to in-progress inputs."""
-
     return _ContainerIOManager.get_input_concurrency()
 
 
@@ -33,7 +31,6 @@ def set_local_input_concurrency(concurrency: int):
     """Set the container's local input concurrency. Dynamic concurrency will be disabled.
     When setting to a smaller value, this method will not interrupt in-progress inputs.
     """
-
     _ContainerIOManager.set_input_concurrency(concurrency)
 
 

modal/functions.py

@@ -299,13 +299,12 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     """Functions are the basic units of serverless execution on Modal.
 
     Generally, you will not construct a `Function` directly. Instead, use the
-    `@app.function()` decorator on the `App` object (formerly called "Stub")
-    for your application.
+    `App.function()` decorator to register your Python functions with your App.
     """
 
     # TODO: more type annotations
     _info: Optional[FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[_Mount]  # set at load time, only by loader
+    _serve_mounts: typing.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]
@@ -315,7 +314,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     _tag: str
     _raw_f: Callable[..., Any]
     _build_args: dict
-    _can_use_base_function: bool = False  # whether we need to call FunctionBindParams
+
     _is_generator: Optional[bool] = None
     _cluster_size: Optional[int] = None
 
@@ -324,10 +323,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     _use_function_id: str  # The function to invoke
     _use_method_name: str = ""
 
-    # TODO (elias): remove _parent. In case of instance functions, and methods bound on those,
-    #  this references the parent class-function and is used to infer the client for lazy-loaded methods
-    _parent: Optional["_Function"] = None
-
     _class_parameter_info: Optional["api_pb2.ClassParameterInfo"] = None
     _method_handle_metadata: Optional[Dict[str, "api_pb2.FunctionHandleMetadata"]] = None
 
@@ -512,7 +507,6 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         fun._info = class_bound_method._info
         fun._obj = instance_service_function._obj
         fun._is_method = True
-        fun._parent = instance_service_function._parent
         fun._app = class_bound_method._app
         fun._spec = class_bound_method._spec
         return fun
@@ -580,6 +574,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
 
         if is_local():
             entrypoint_mounts = info.get_entrypoint_mount()
+
             all_mounts = [
                 _get_client_mount(),
                 *explicit_mounts,
@@ -612,6 +607,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         if proxy:
             # HACK: remove this once we stop using ssh tunnels for this.
             if image:
+                # TODO(elias): this will cause an error if users use prior `.add_local_*` commands without copy=True
                 image = image.apt_install("autossh")
 
         function_spec = _FunctionSpec(
@@ -828,7 +824,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     )
                     for path, volume in validated_volumes
                 ]
-                loaded_mount_ids = {m.object_id for m in all_mounts}
+                loaded_mount_ids = {m.object_id for m in all_mounts} | {m.object_id for m in image._mount_layers}
 
                 # Get object dependencies
                 object_dependencies = []
@@ -970,9 +966,9 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         raise InvalidError(f"Function {info.function_name} is too large to deploy.")
                     raise
                 function_creation_status.set_response(response)
-            local_mounts = set(m for m in all_mounts if m.is_local())  # needed for modal.serve file watching
-            local_mounts |= image._used_local_mounts
-            obj._used_local_mounts = frozenset(local_mounts)
+            serve_mounts = set(m for m in all_mounts if m.is_local())  # needed for modal.serve file watching
+            serve_mounts |= image._serve_mounts
+            obj._serve_mounts = frozenset(serve_mounts)
             self._hydrate(response.function_id, resolver.client, response.handle_metadata)
 
         rep = f"Function({tag})"
@@ -1018,27 +1014,37 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         Binds a class-function to a specific instance of (init params, options) or a new workspace
         """
 
-        async def _load(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
-            if self._parent is None:
+        # In some cases, reuse the base function, i.e. not create new clones of each method or the "service function"
+        can_use_parent = len(args) + len(kwargs) == 0 and not from_other_workspace and options is None
+        parent = self
+
+        async def _load(param_bound_func: _Function, resolver: Resolver, existing_object_id: Optional[str]):
+            if parent is None:
                 raise ExecutionError("Can't find the parent class' service function")
             try:
-                identity = f"{self._parent.info.function_name} class service function"
+                identity = f"{parent.info.function_name} class service function"
             except Exception:
                 # Can't always look up the function name that way, so fall back to generic message
                 identity = "class service function for a parameterized class"
-            if not self._parent.is_hydrated:
-                if self._parent.app._running_app is None:
-                    reason = ", because the App it is defined on is not running."
+            if not parent.is_hydrated:
+                if parent.app._running_app is None:
+                    reason = ", because the App it is defined on is not running"
                 else:
                     reason = ""
                 raise ExecutionError(
                     f"The {identity} has not been hydrated with the metadata it needs to run on Modal{reason}."
                 )
-            assert self._parent._client.stub
+
+            assert 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.
+                param_bound_func._hydrate_from_other(parent)
+                return
+
             if (
-                self._parent._class_parameter_info
-                and self._parent._class_parameter_info.format
-                == api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PROTO
+                parent._class_parameter_info
+                and parent._class_parameter_info.format == api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PROTO
             ):
                 if args:
                     # TODO(elias) - We could potentially support positional args as well, if we want to?
@@ -1046,34 +1052,30 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         "Can't use positional arguments with modal.parameter-based synthetic constructors.\n"
                         "Use (<parameter_name>=value) keyword arguments when constructing classes instead."
                     )
-                serialized_params = serialize_proto_params(kwargs, self._parent._class_parameter_info.schema)
+                serialized_params = serialize_proto_params(kwargs, parent._class_parameter_info.schema)
             else:
                 serialized_params = serialize((args, kwargs))
             environment_name = _get_environment_name(None, resolver)
-            assert self._parent is not None
+            assert parent is not None
             req = api_pb2.FunctionBindParamsRequest(
-                function_id=self._parent._object_id,
+                function_id=parent._object_id,
                 serialized_params=serialized_params,
                 function_options=options,
                 environment_name=environment_name
                 or "",  # TODO: investigate shouldn't environment name always be specified here?
             )
 
-            response = await retry_transient_errors(self._parent._client.stub.FunctionBindParams, req)
-            self._hydrate(response.bound_function_id, self._parent._client, response.handle_metadata)
+            response = await retry_transient_errors(parent._client.stub.FunctionBindParams, req)
+            param_bound_func._hydrate(response.bound_function_id, parent._client, response.handle_metadata)
 
         fun: _Function = _Function._from_loader(_load, "Function(parametrized)", hydrate_lazily=True)
 
-        # In some cases, reuse the base function, i.e. not create new clones of each method or the "service function"
-        fun._can_use_base_function = len(args) + len(kwargs) == 0 and not from_other_workspace and options is None
-        if fun._can_use_base_function and self.is_hydrated:
-            # Edge case that lets us hydrate all objects right away
-            # if the instance didn't use explicit constructor arguments
-            fun._hydrate_from_other(self)
+        if can_use_parent and parent.is_hydrated:
+            # skip the resolver altogether:
+            fun._hydrate_from_other(parent)
 
         fun._info = self._info
         fun._obj = obj
-        fun._parent = self
         return fun
 
     @live_method
@@ -1223,7 +1225,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         self._function_name = None
         self._info = None
         self._use_function_id = ""
-        self._used_local_mounts = frozenset()
+        self._serve_mounts = frozenset()
 
     def _hydrate_metadata(self, metadata: Optional[Message]):
         # Overridden concrete implementation of base class method
@@ -1264,8 +1266,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 + f"or call it locally: {self._function_name}.local()"
             )
 
+    # TODO (live_method on properties is not great, since it could be blocking the event loop from async contexts)
     @property
-    def web_url(self) -> str:
+    @live_method
+    async def web_url(self) -> str:
         """URL of a Function running as a web endpoint."""
         if not self._web_url:
             raise ValueError(
@@ -1438,7 +1442,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             return fun(*args, **kwargs)
         else:
             # This is a method on a class, so bind the self to the function
-            user_cls_instance = obj._get_user_cls_instance()
+            user_cls_instance = obj._cached_user_cls_instance()
 
             fun = info.raw_f.__get__(user_cls_instance)
 

modal/functions.pyi

@@ -110,7 +110,7 @@ OriginalReturnType = typing.TypeVar("OriginalReturnType", covariant=True)
 
 class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object._Object):
     _info: typing.Optional[modal._utils.function_utils.FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[modal.mount._Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount._Mount]
     _app: typing.Optional[modal.app._App]
     _obj: typing.Optional[modal.cls._Obj]
     _web_url: typing.Optional[str]
@@ -120,12 +120,10 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
     _tag: str
     _raw_f: typing.Callable[..., typing.Any]
     _build_args: dict
-    _can_use_base_function: bool
     _is_generator: typing.Optional[bool]
     _cluster_size: typing.Optional[int]
     _use_function_id: str
     _use_method_name: str
-    _parent: typing.Optional[_Function]
     _class_parameter_info: typing.Optional[modal_proto.api_pb2.ClassParameterInfo]
     _method_handle_metadata: typing.Optional[typing.Dict[str, modal_proto.api_pb2.FunctionHandleMetadata]]
 
@@ -218,7 +216,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.
     def _get_metadata(self): ...
     def _check_no_web_url(self, fn_name: str): ...
     @property
-    def web_url(self) -> str: ...
+    async def web_url(self) -> str: ...
     @property
     def is_generator(self) -> bool: ...
     @property
@@ -286,7 +284,7 @@ P_INNER = typing_extensions.ParamSpec("P_INNER")
 
 class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.Object):
     _info: typing.Optional[modal._utils.function_utils.FunctionInfo]
-    _used_local_mounts: typing.FrozenSet[modal.mount.Mount]
+    _serve_mounts: typing.FrozenSet[modal.mount.Mount]
     _app: typing.Optional[modal.app.App]
     _obj: typing.Optional[modal.cls.Obj]
     _web_url: typing.Optional[str]
@@ -296,12 +294,10 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     _tag: str
     _raw_f: typing.Callable[..., typing.Any]
     _build_args: dict
-    _can_use_base_function: bool
     _is_generator: typing.Optional[bool]
     _cluster_size: typing.Optional[int]
     _use_function_id: str
     _use_method_name: str
-    _parent: typing.Optional[Function]
     _class_parameter_info: typing.Optional[modal_proto.api_pb2.ClassParameterInfo]
     _method_handle_metadata: typing.Optional[typing.Dict[str, modal_proto.api_pb2.FunctionHandleMetadata]]
 
@@ -450,11 +446,11 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
 
     _call_generator_nowait: ___call_generator_nowait_spec
 
-    class __remote_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class __remote_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
         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]
 
     class __remote_gen_spec(typing_extensions.Protocol):
         def __call__(self, *args, **kwargs) -> typing.Generator[typing.Any, None, None]: ...
@@ -466,17 +462,17 @@ class Function(typing.Generic[P, ReturnType, OriginalReturnType], modal.object.O
     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]):
         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]
 
-    class __spawn_spec(typing_extensions.Protocol[P_INNER, ReturnType_INNER]):
+    class __spawn_spec(typing_extensions.Protocol[ReturnType_INNER, P_INNER]):
         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]
 
     def get_raw_f(self) -> typing.Callable[..., typing.Any]: ...
 

modal/gpu.py

@@ -1,6 +1,6 @@
 # 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
 
@@ -147,25 +147,27 @@ 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,
     "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).
+[`@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.
 """