[Tune] Keep resource specifications when nesting with_resources in with_parameters

[justinvyu]

Why are these changes needed?

Order of nesting matters

The order of nesting matters when using tune.with_parameters and tune.with_resources together. When tune.with_resources is nested inside a tune.with_parameters, the resource specification is dropped and not considered when launching the trials.

def train_func(config, extra_data=None):
    pass

# This works (resources assigned correctly)
trainable = tune.with_resources(
    tune.with_parameters(train_func, extra_data=1),
    {"cpu": 1.0, "gpu": 1.0}
)
results = Tuner(trainable).fit()

# This doesn't work!
trainable = tune.with_parameters(
    tune.with_resources(train_func, {"cpu": 1.0, "gpu": 1.0}),
    extra_data=1
)
results = Tuner(trainable).fit()

This PR introduces a fix for both cases above to work.

tune.with_parameters on a training function with a return value

Another issue addressed in this PR (let me know if this should be its own PR) is allowing tune.with_parameters to work with custom training functions that return a value.

def train_func(config, extra_data=None):
    return {"metric": 1}  # train function returns a "final result"

# The return value gets dropped after wrapping with `tune.with_parameters`
trainable = tune.with_parameters(train_func, extra_data=1)

This PR returns the original training function's return value instead of dropping it.

Open questions

1. If we introduce any more trainable wrapper utils, we will need to maintain compatibility by passing around the _resource attribute. Is there a better way to keep the resource specification information around?

Related issue number

Closes #29235

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[krfricke]

Great, thanks!