[tune] Support tune.with_resources for class methods

python/ray/tune/tests/test_api.py

@@ -1494,6 +1494,27 @@ def default_resource_request(cls, config):
     assert trial.last_result["_metric"] == num_gpus
 
 
+@pytest.mark.parametrize("num_gpus", [1, 2])
+def test_with_resources_class_method(ray_start_2_cpus_2_gpus, num_gpus):
+    class Worker:
+        def train_fn(self, config):
+            return len(ray.get_gpu_ids())
+
+    worker = Worker()
+
+    [trial] = tune.run(
+        tune.with_resources(
+            worker.train_fn,
+            resources=lambda config: PlacementGroupFactory(
+                [{"GPU": config["use_gpus"]}]
+            ),
+        ),
+        config={"use_gpus": num_gpus},
+    ).trials
+
+    assert trial.last_result["_metric"] == num_gpus
+
+
 class SerializabilityTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):

python/ray/tune/trainable/util.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import shutil
+import types
 from typing import Any, Callable, Dict, Optional, Type, Union, TYPE_CHECKING
 
 import pandas as pd
@@ -443,8 +444,31 @@ def train(config):
         )
 
     if not inspect.isclass(trainable):
+        if isinstance(trainable, types.MethodType):
+            # Methods cannot set arbitrary attributes, so we have to wrap them
+            use_checkpoint = _detect_checkpoint_function(trainable, partial=True)
+            if use_checkpoint:
+
+                def _trainable(config, checkpoint_dir):
+                    return trainable(config, checkpoint_dir=checkpoint_dir)
+
+            else:
+
+                def _trainable(config):
+                    return trainable(config)
+
+            _trainable._resources = pgf
+            return _trainable
+
         # Just set an attribute. This will be resolved later in `wrap_function()`.
-        trainable._resources = pgf
+        try:
+            trainable._resources = pgf
+        except AttributeError as e:
+            raise RuntimeError(
+                "Could not use `tune.with_resources()` on the supplied trainable. "
+                "Wrap your trainable in a regular function before passing it "
+                "to Ray Tune."
+            ) from e
     else:
 
         class ResourceTrainable(trainable):