[1/n] Lightweight Ray AIR API refactor

[pcmoritz]

Why are these changes needed?

This PR removes some circularities in the Ray AIR import system so we can put the training related functions into ray.train. It introduces a training context and makes report, get_dataset_shard, get_context, Checkpoint, Result, and the following configs:

• CheckpointConfig
• DataConfig
• FailureConfig
• RunConfig
• ScalingConfig

available in ray.train. No user facing changes yet, the old APIs still work.

Going forward, it will be most consistent / symmetrical if these things are included in the following way:

from ray import train, tune, serve # Pick the subset that is needed
# Include what you need from the following:
from ray.train import CheckpointConfig, DataConfig, FailureConfig, RunConfig, ScalingConfig

# ...

def train_func():
    dataset_shard = train.get_dataset_shard("train")
    world_size = train.get_context().get_world_size()
    # ...
    train.report(...)

trainer = train.torch.TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=2),
)
result = trainer.fit()

We have many examples in #37123 on how this looks like in actual code.

Related issue number

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

[ericl]

Nice. Is this actually the bulk of the changes (aside from checkpoint class)? I guess I can think of the following followups:

• Make it so you can run any trainable func/class with Trainer.
• Update docs/deprecate old APIs.

[pcmoritz]

This is what I have seen so far from Goku's course. There is also the Result type that is currently still in AIR. Make it so you can run any trainable func/class with Trainer. <- let's chat more about this, I'm actually not super sure what needs to be done here.

After this is merged, I will start working on migrating our examples / deprecation, but I don't expect it to land until after the 2.6 branch cut :)

[ericl]

Do we want to avoid moving the Checkpoint class if we likely will remove it?

[pcmoritz]

Alright, I will remove them.

[ericl]

Should these accessors be part of session only?

[pcmoritz]

I want to get rid of the session namespace going forward. It serves no purpose.

[ericl]

I feel like we should do that atomically then, it seems odd to move just a few specific methods initially. It is also a big change and I'm not clear if the pros outweigh the cons.

[pcmoritz]

So I don't think get_local_rank and get_local_world_size are currently being used anywhere (they are internal APIs and used to automatically set an environment variable for TorchTrainer). So I'd like to get rid of them. Should I be wrong I'm happy to add them later. So is get_node_rank (which is also very confusing btw). Btw, the docs strings for all of these except get_local_rank is wrong too.

For get_trial_dir, get_trial_id, get_trial_name, get_experiment_name I would first like to understand how your checkpointing changes are changing things, but I'm happy to add them back if you think they will keep being the same.

get_checkpoint will probably become irrelevant.

That's already all the methods in session.

[ericl]

So these are all important utility methods. The more general concern is that there will always be more and more of these small things, just see ray.get_runtime_context() for example (I believe session should be renamed something like context, IMO).

We should plan for a large number of utility methods existing here in the future, and I don't think it's great to have them directly scoped under train.

[pcmoritz]

Yeah, my ideal API for this would be

def train_loop_per_worker(context):
    # Use context.report(...), context.get_dataset_shard(...) etc.

Let me look into whether this can be made to happen in a compatible way.

[matthewdeng]

+1 to renaming Session to *Context.

I don't really see a benefit of passing it into the train_loop_per_worker or trainable, so I would prefer not to unless there is a strong enough reason for this.

[pcmoritz]

Alright, the plan changed to renaming from ray.air import session -> from ray.train import context :)

[matthewdeng]

This is a Tune concept.

[ericl]

We're moving this as well to Train right? You cannot launch trials without resources.

[pcmoritz]

Yes, this is a train concept and being used in train examples (e.g. to set the number of threads for trainers).

[matthewdeng]

Maybe we can discuss this along with:

Make it so you can run any trainable func/class with Trainer.

Right now the interfaces are:
ScalingConfig (Train) -> PlacementGroupFactory (Tune) -> PlacementGroup (Core)

I would say that the usage in Train examples is a misuse of the API and extremely confusing. For setting the number of threads, this is being used as a proxy when we actually just want an API to get the number of CPUs assigned to the RayTrainWorker actor.

[ericl]

Yea the placement group stuff also needs to move in this case.