Remove usage of method placeholder functions

modal/_runtime/container_io_manager.py

@@ -452,7 +452,7 @@ async def _dynamic_concurrency_loop(self):
             await asyncio.sleep(DYNAMIC_CONCURRENCY_INTERVAL_SECS)
 
     async def get_app_objects(self) -> RunningApp:
-        req = api_pb2.AppGetObjectsRequest(app_id=self.app_id, include_unindexed=True)
+        req = api_pb2.AppGetObjectsRequest(app_id=self.app_id, include_unindexed=True, only_class_function=True)
         resp = await retry_transient_errors(self._client.stub.AppGetObjects, req)
         logger.debug(f"AppGetObjects received {len(resp.items)} objects for app {self.app_id}")
 

modal/cli/import_refs.py

@@ -110,7 +110,7 @@ def get_by_object_path(obj: Any, obj_path: str) -> Optional[Any]:
 def _infer_function_or_help(
     app: App, module, accept_local_entrypoint: bool, accept_webhook: bool
 ) -> Union[Function, LocalEntrypoint]:
-    function_choices = set(tag for tag, func in app.registered_functions.items() if not func.info.is_service_class())
+    function_choices = set(tag for tag in app.registered_functions.keys())
     if not accept_webhook:
         function_choices -= set(app.registered_web_endpoints)
     if accept_local_entrypoint:

modal/cli/run.py

@@ -136,20 +136,33 @@ def _get_clean_app_description(func_ref: str) -> str:
 
 
 def _get_click_command_for_function(app: App, function_tag):
-    function = app.registered_functions[function_tag]
+    function = app.registered_functions.get(function_tag, None)
+    if not function or (isinstance(function, Function) and function.info.user_cls is not None):
+        # This is either a function_tag for a class method function (e.g MyClass.foo) or a function tag for a
+        # class service function (MyClass.*)
+        class_name, method_name = function_tag.rsplit(".", 1)
+        if not function:
+            function = app.registered_functions.get(f"{class_name}.*")
     assert isinstance(function, Function)
     function = typing.cast(Function, function)
     if function.is_generator:
         raise InvalidError("`modal run` is not supported for generator functions")
 
     signature: Dict[str, ParameterMetadata]
     cls: Optional[Cls] = None
-    method_name: Optional[str] = None
     if function.info.user_cls is not None:
-        class_name, method_name = function_tag.rsplit(".", 1)
         cls = typing.cast(Cls, app.registered_classes[class_name])
         cls_signature = _get_signature(function.info.user_cls)
-        fun_signature = _get_signature(function.info.raw_f, is_method=True)
+        if method_name == "*":
+            method_names = list(cls._get_partial_functions().keys())
+            if len(method_names) == 1:
+                method_name = method_names[0]
+            else:
+                class_name = function.info.user_cls.__name__
+                raise click.UsageError(
+                    f"Please specify a specific method of {class_name} to run, e.g. `modal run foo.py::MyClass.bar`"  # noqa: E501
+                )
+        fun_signature = _get_signature(getattr(cls, method_name).info.raw_f, is_method=True)
         signature = dict(**cls_signature, **fun_signature)  # Pool all arguments
         # TODO(erikbern): assert there's no overlap?
     else:

modal/cls.py

@@ -244,7 +244,7 @@ class _Cls(_Object, type_prefix="cs"):
     _class_service_function: Optional[
         _Function
     ]  # The _Function serving *all* methods of the class, used for version >=v0.63
-    _method_functions: Dict[str, _Function]  # Placeholder _Functions for each method
+    _method_functions: Dict[str, _Function] = None  # Placeholder _Functions for each method
     _options: Optional[api_pb2.FunctionOptions]
     _callables: Dict[str, Callable[..., Any]]
     _from_other_workspace: Optional[bool]  # Functions require FunctionBindParams before invocation.
@@ -253,7 +253,6 @@ class _Cls(_Object, type_prefix="cs"):
     def _initialize_from_empty(self):
         self._user_cls = None
         self._class_service_function = None
-        self._method_functions = {}
         self._options = None
         self._callables = {}
         self._from_other_workspace = None
@@ -273,27 +272,49 @@ def _get_partial_functions(self) -> Dict[str, _PartialFunction]:
 
     def _hydrate_metadata(self, metadata: Message):
         assert isinstance(metadata, api_pb2.ClassHandleMetadata)
-
-        for method in metadata.methods:
-            if method.function_name in self._method_functions:
-                # This happens when the class is loaded locally
-                # since each function will already be a loaded dependency _Function
-                self._method_functions[method.function_name]._hydrate(
-                    method.function_id, self._client, method.function_handle_metadata
-                )
+        if (
+            self._class_service_function
+            and self._class_service_function._method_handle_metadata
+            and len(self._class_service_function._method_handle_metadata)
+        ):
+            # The class only has a class service service function and no method placeholders.
+            if self._method_functions:
+                # We're here when the Cls is loaded locally (e.g. _Cls.from_local) so the _method_functions mapping is
+                # populated with (un-hydrated) _Function objects
+                for (
+                    method_name,
+                    method_handle_metadata,
+                ) in self._class_service_function._method_handle_metadata.items():
+                    self._method_functions[method_name]._hydrate(
+                        self._class_service_function.object_id, self._client, method_handle_metadata
+                    )
             else:
+                # We're here when the function is loaded remotely (e.g. _Cls.from_name)
+                self._method_functions = {}
+                for (
+                    method_name,
+                    method_handle_metadata,
+                ) in self._class_service_function._method_handle_metadata.items():
+                    self._method_functions[method_name] = _Function._new_hydrated(
+                        self._class_service_function.object_id, self._client, method_handle_metadata
+                    )
+        elif self._class_service_function:
+            # A class with a class service function and method placeholder functions
+            self._method_functions = {}
+            for method in metadata.methods:
+                self._method_functions[method.function_name] = _Function._new_hydrated(
+                    self._class_service_function.object_id, self._client, method.function_handle_metadata
+                )
+        else:
+            # pre 0.63 class that does not have a class service function and only method functions
+            self._method_functions = {}
+            for method in metadata.methods:
                 self._method_functions[method.function_name] = _Function._new_hydrated(
                     method.function_id, self._client, method.function_handle_metadata
                 )
 
     def _get_metadata(self) -> api_pb2.ClassHandleMetadata:
         class_handle_metadata = api_pb2.ClassHandleMetadata()
-        for f_name, f in self._method_functions.items():
-            class_handle_metadata.methods.append(
-                api_pb2.ClassMethod(
-                    function_name=f_name, function_id=f.object_id, function_handle_metadata=f._get_metadata()
-                )
-            )
         return class_handle_metadata
 
     @staticmethod
@@ -327,56 +348,43 @@ def from_local(user_cls, app: "modal.app._App", class_service_function: _Functio
         # validate signature
         _Cls.validate_construction_mechanism(user_cls)
 
-        functions: Dict[str, _Function] = {}
+        method_functions: Dict[str, _Function] = {}
         partial_functions: Dict[str, _PartialFunction] = _find_partial_methods_for_user_cls(
             user_cls, _PartialFunctionFlags.FUNCTION
         )
 
         for method_name, partial_function in partial_functions.items():
-            method_function = class_service_function._bind_method_old(user_cls, method_name, partial_function)
-            app._add_function(method_function, is_web_endpoint=partial_function.webhook_config is not None)
+            method_function = class_service_function._bind_method(user_cls, method_name, partial_function)
+            if partial_function.webhook_config is not None:
+                app._web_endpoints.append(method_function.tag)
             partial_function.wrapped = True
-            functions[method_name] = method_function
+            method_functions[method_name] = method_function
 
         # Disable the warning that these are not wrapped
         for partial_function in _find_partial_methods_for_user_cls(user_cls, ~_PartialFunctionFlags.FUNCTION).values():
             partial_function.wrapped = True
 
         # Get all callables
         callables: Dict[str, Callable] = {
-            k: pf.raw_f for k, pf in _find_partial_methods_for_user_cls(user_cls, ~_PartialFunctionFlags(0)).items()
+            k: pf.raw_f for k, pf in _find_partial_methods_for_user_cls(user_cls, _PartialFunctionFlags.all()).items()
         }
 
         def _deps() -> List[_Function]:
-            return [class_service_function] + list(functions.values())
+            return [class_service_function]
 
         async def _load(self: "_Cls", resolver: Resolver, existing_object_id: Optional[str]):
-            req = api_pb2.ClassCreateRequest(app_id=resolver.app_id, existing_class_id=existing_object_id)
-            for f_name, f in self._method_functions.items():
-                req.methods.append(
-                    api_pb2.ClassMethod(
-                        function_name=f_name, function_id=f.object_id, function_handle_metadata=f._get_metadata()
-                    )
-                )
+            req = api_pb2.ClassCreateRequest(
+                app_id=resolver.app_id, existing_class_id=existing_object_id, only_class_function=True
+            )
             resp = await resolver.client.stub.ClassCreate(req)
-            # Even though we already have the function_handle_metadata for this method locally,
-            # The RPC is going to replace it with function_handle_metadata derived from the server.
-            # We need to overwrite the definition_id sent back from the server here with the definition_id
-            # previously stored in function metadata, which may have been sent back from FunctionCreate.
-            # The problem is that this metadata propagates back and overwrites the metadata on the Function
-            # object itself. This is really messy. Maybe better to exclusively populate the method metadata
-            # from the function metadata we already have locally? Really a lot to clean up here...
-            for method in resp.handle_metadata.methods:
-                f_metadata = self._method_functions[method.function_name]._get_metadata()
-                method.function_handle_metadata.definition_id = f_metadata.definition_id
             self._hydrate(resp.class_id, resolver.client, resp.handle_metadata)
 
         rep = f"Cls({user_cls.__name__})"
         cls: _Cls = _Cls._from_loader(_load, rep, deps=_deps)
         cls._app = app
         cls._user_cls = user_cls
         cls._class_service_function = class_service_function
-        cls._method_functions = functions
+        cls._method_functions = method_functions
         cls._callables = callables
         cls._from_other_workspace = False
         return cls
@@ -415,6 +423,7 @@ async def _load_remote(obj: _Object, resolver: Resolver, existing_object_id: Opt
                 environment_name=_environment_name,
                 lookup_published=workspace is not None,
                 workspace_name=workspace,
+                only_class_function=True,
             )
             try:
                 response = await retry_transient_errors(resolver.client.stub.ClassGet, request)

modal/functions.py

@@ -118,7 +118,7 @@ async def create(
         function_call_invocation_type: "api_pb2.FunctionCallInvocationType.ValueType",
     ) -> "_Invocation":
         assert client.stub
-        function_id = function._invocation_function_id()
+        function_id = function.object_id
         item = await _create_input(args, kwargs, client, method_name=function._use_method_name)
 
         request = api_pb2.FunctionMapRequest(
@@ -319,8 +319,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
     _cluster_size: Optional[int] = None
 
     # when this is the method of a class/object function, invocation of this function
-    # should be using another function id and supply the method name in the FunctionInput:
-    _use_function_id: str  # The function to invoke
+    # should supply the method name in the FunctionInput:
     _use_method_name: str = ""
 
     _class_parameter_info: Optional["api_pb2.ClassParameterInfo"] = None
@@ -360,94 +359,6 @@ def _bind_method(
         fun._is_method = True
         return fun
 
-    def _bind_method_old(
-        self,
-        user_cls,
-        method_name: str,
-        partial_function: "modal.partial_function._PartialFunction",
-    ):
-        """mdmd:hidden
-
-        Creates a function placeholder function that binds a specific method name to
-        this function for use when invoking the function.
-
-        Should only be used on "class service functions". For "instance service functions",
-        we don't create an actual backend function, and instead do client-side "fake-hydration"
-        only, see _bind_instance_method.
-
-        """
-        class_service_function = self
-        assert class_service_function._info  # has to be a local function to be able to "bind" it
-        assert not class_service_function._is_method  # should not be used on an already bound method placeholder
-        assert not class_service_function._obj  # should only be used on base function / class service function
-        full_name = f"{user_cls.__name__}.{method_name}"
-        function_type = get_function_type(partial_function.is_generator)
-
-        async def _load(method_bound_function: "_Function", resolver: Resolver, existing_object_id: Optional[str]):
-            function_definition = api_pb2.Function(
-                function_name=full_name,
-                webhook_config=partial_function.webhook_config,
-                function_type=function_type,
-                is_method=True,
-                use_function_id=class_service_function.object_id,
-                use_method_name=method_name,
-                batch_max_size=partial_function.batch_max_size or 0,
-                batch_linger_ms=partial_function.batch_wait_ms or 0,
-            )
-            assert resolver.app_id
-            request = api_pb2.FunctionCreateRequest(
-                app_id=resolver.app_id,
-                function=function_definition,
-                #  method_bound_function.object_id usually gets set by preload
-                existing_function_id=existing_object_id or method_bound_function.object_id or "",
-                defer_updates=True,
-            )
-            assert resolver.client.stub is not None  # client should be connected when load is called
-            with FunctionCreationStatus(resolver, full_name) as function_creation_status:
-                response = await resolver.client.stub.FunctionCreate(request)
-                method_bound_function._hydrate(
-                    response.function_id,
-                    resolver.client,
-                    response.handle_metadata,
-                )
-                function_creation_status.set_response(response)
-
-        async def _preload(method_bound_function: "_Function", resolver: Resolver, existing_object_id: Optional[str]):
-            if class_service_function._use_method_name:
-                raise ExecutionError(f"Can't bind method to already bound {class_service_function}")
-            assert resolver.app_id
-            req = api_pb2.FunctionPrecreateRequest(
-                app_id=resolver.app_id,
-                function_name=full_name,
-                function_type=function_type,
-                webhook_config=partial_function.webhook_config,
-                use_function_id=class_service_function.object_id,
-                use_method_name=method_name,
-                existing_function_id=existing_object_id or "",
-            )
-            assert resolver.client.stub  # client should be connected at this point
-            response = await retry_transient_errors(resolver.client.stub.FunctionPrecreate, req)
-            method_bound_function._hydrate(response.function_id, resolver.client, response.handle_metadata)
-
-        def _deps():
-            return [class_service_function]
-
-        rep = f"Method({full_name})"
-
-        fun = _Function._from_loader(_load, rep, preload=_preload, deps=_deps)
-        fun._tag = full_name
-        fun._raw_f = partial_function.raw_f
-        fun._info = FunctionInfo(
-            partial_function.raw_f, user_cls=user_cls, serialized=class_service_function.info.is_serialized()
-        )  # needed for .local()
-        fun._use_method_name = method_name
-        fun._app = class_service_function._app
-        fun._is_generator = partial_function.is_generator
-        fun._cluster_size = partial_function.cluster_size
-        fun._spec = class_service_function._spec
-        fun._is_method = True
-        return fun
-
     def _bind_instance_method(self, class_bound_method: "_Function"):
         """mdmd:hidden
 
@@ -475,7 +386,6 @@ def hydrate_from_instance_service_function(method_placeholder_fun):
             method_placeholder_fun._is_generator = class_bound_method._is_generator
             method_placeholder_fun._cluster_size = class_bound_method._cluster_size
             method_placeholder_fun._use_method_name = method_name
-            method_placeholder_fun._use_function_id = instance_service_function.object_id
             method_placeholder_fun._is_method = True
 
         async def _load(fun: "_Function", resolver: Resolver, existing_object_id: Optional[str]):
@@ -848,6 +758,8 @@ async def _load(self: _Function, resolver: Resolver, existing_object_id: Optiona
                     class_serialized=class_serialized or b"",
                     function_type=function_type,
                     webhook_config=webhook_config,
+                    method_definitions=method_definitions,
+                    method_definitions_set=True,
                     shared_volume_mounts=network_file_system_mount_protos(
                         validated_network_file_systems, allow_cross_region_volumes
                     ),
@@ -1224,7 +1136,6 @@ def _initialize_from_empty(self):
         self._web_url = None
         self._function_name = None
         self._info = None
-        self._use_function_id = ""
         self._serve_mounts = frozenset()
 
     def _hydrate_metadata(self, metadata: Optional[Message]):
@@ -1234,15 +1145,11 @@ def _hydrate_metadata(self, metadata: Optional[Message]):
         self._web_url = metadata.web_url
         self._function_name = metadata.function_name
         self._is_method = metadata.is_method
-        self._use_function_id = metadata.use_function_id
         self._use_method_name = metadata.use_method_name
         self._class_parameter_info = metadata.class_parameter_info
         self._method_handle_metadata = dict(metadata.method_handle_metadata)
         self._definition_id = metadata.definition_id
 
-    def _invocation_function_id(self) -> str:
-        return self._use_function_id or self.object_id
-
     def _get_metadata(self):
         # Overridden concrete implementation of base class method
         assert self._function_name, f"Function name must be set before metadata can be retrieved for {self}"
@@ -1251,7 +1158,6 @@ def _get_metadata(self):
             function_type=get_function_type(self._is_generator),
             web_url=self._web_url or "",
             use_method_name=self._use_method_name,
-            use_function_id=self._use_function_id,
             is_method=self._is_method,
             class_parameter_info=self._class_parameter_info,
             definition_id=self._definition_id,