[Data] Propagate driver DataContext to RayTrainWorkers

[scottjlee]

Why are these changes needed?

Second attempt on #39698, which was found to be incompatible with RLLib Learner classes. In this PR, we instead move the logic of passing the driver's DataContext into the BackendExecutor, instead of the RayTrainWorker as previously.

Related issue number

Closes #39237
Previous PR: #39698

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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[scottjlee]

CI run with ML / RL tests passing: https://buildkite.com/ray-project/oss-ci-build-pr/builds/37993

Going to now revert manual enabling the RL tests trigger.

[woshiyyya]

Thanks @scottjlee, this solution is cleaner than the previous one.

Also, can you elaborate more on the RLLib Learner issue?

[scottjlee]

Thanks @scottjlee, this solution is cleaner than the previous one.

Also, can you elaborate more on the RLLib Learner issue?

Yeah, the previous implementation, which added a new parameter into RayTrainWorker.__init__(), was incompatible with Learner which doesn't have this extra parameter (and we do not want to change this RLLib API).

[matthewdeng]

Is this intentional as part of this PR?

[scottjlee]

yeah, paired with @justinvyu on this for some time, and we came to the conclusion that the mocking inside the test may need to be updated to be compatible with the fix in this PR, but we couldn't figure it out. I think @justinvyu said he can come back in the future to fix or remove the test, will also let him elaborate