[Tune] Keep resource specifications when nesting with_resources in with_parameters

python/ray/tune/tests/test_api.py

@@ -1530,6 +1530,34 @@ def train_fn(self, config):
     assert trial.last_result["_metric"] == num_gpus
 
 
+@pytest.mark.parametrize("num_gpus", [1, 2])
+def test_with_resources_and_parameters_fn(ray_start_2_cpus_2_gpus, num_gpus):
+    def train_fn(config, extra_param=None):
+        assert extra_param is not None, "Missing extra parameter."
+        print(ray.get_runtime_context().get_assigned_resources())
+        return {"num_gpus": len(ray.get_gpu_ids())}
+
+    # Nesting `tune.with_parameters` and `tune.with_resources` should respect
+    # the resource specifications.
+    trainable = tune.with_resources(
+        tune.with_parameters(train_fn, extra_param="extra"),
+        {"gpu": num_gpus},
+    )
+
+    tuner = tune.Tuner(trainable)
+    results = tuner.fit()
+    print(results[0].metrics)
+    assert results[0].metrics["num_gpus"] == num_gpus
+
+    # The other order of nesting should work the same.
+    trainable = tune.with_parameters(
+        tune.with_resources(train_fn, {"gpu": num_gpus}), extra_param="extra"
+    )
+    tuner = tune.Tuner(trainable)
+    results = tuner.fit()
+    assert results[0].metrics["num_gpus"] == num_gpus
+
+
 class SerializabilityTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):

python/ray/tune/trainable/util.py

@@ -362,19 +362,28 @@ def inner(config, checkpoint_dir=None):
 
             for k in keys:
                 fn_kwargs[k] = parameter_registry.get(prefix + k)
-            trainable(config, **fn_kwargs)
+            return trainable(config, **fn_kwargs)
 
         inner.__name__ = trainable_name
 
+        # If the trainable has been wrapped with `tune.with_resources`, we should
+        # keep the `_resources` attribute around
+        if hasattr(trainable, "_resources"):
+            inner._resources = trainable._resources
+
         # Use correct function signature if no `checkpoint_dir` parameter
         # is set
         if not use_checkpoint:
 
             def _inner(config):
-                inner(config, checkpoint_dir=None)
+                return inner(config, checkpoint_dir=None)
 
             _inner.__name__ = trainable_name
 
+            # Again, pass along the resource specification if it exists
+            if hasattr(inner, "_resources"):
+                _inner._resources = inner._resources
+
             if hasattr(trainable, "__mixins__"):
                 _inner.__mixins__ = trainable.__mixins__
             return _inner