[Data] Partial fix for Dataset.context not being sealed after creation

[raulchen]

Why are these changes needed?

Dataset.context should be sealed the first time the Dataset is created. But if a new operator is applied to the dataset, the new global DataContext will be saved again to the Dataset.

This bug prevents using different DataContexts for training and validation datasets in a training job.

Note this PR only fixes the issue when multiple datasets are created in the process but will be running in different processes. If they run in the same process, it's still a bug, see #41573.

Related issue number

#41573

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 :(

[raulchen]

Moving this test to test_context_propagation, not changing the test code.

[stephanie-wang]

What about the places where we use DataContext.get_current() during planning, e.g., here? Don't we need to propagate the DataContext through to those?

[raulchen]

What about the places where we use DataContext.get_current() during planning, e.g., here? Don't we need to propagate the DataContext through to those?

Good point. So this PR can fix the case for training jobs, where be different datasets will be proposed to different processes (the SplitCoordinator actors) for execution.
But if they run in the driver process, it's still a bug. I created an issue #41573 for tracking.

[stephanie-wang]

Looks good, but can you update the PR description to make it clear what cases this does and does not cover?