[REP] Refining the Ray AIR Surface API

[pcmoritz]

This REP proposes to remove the ray.air namespace and put the functionality into the respective libraries Ray Data, Ray Serve and Ray Train.

[ericl]

Could we flesh out an example of this migration in the REP? It may make it more clear the implications of this change.

[pcmoritz]

Yes I will do that, great point.

[pcmoritz]

added now

[krfricke]

I'm fine with this proposal and only have detail questions left.

Just for completeness (as it's not mentioned in the REP). An alternative to resolving this is by moving everything into AIR instead. Thus we would have ray.air.training, ray.air.tuning, etc. In such a world we would soften the boundaries between the libraries and keep shared modules in the AIR namespace. By moving the libraries "one level down" we would enforce a separation from Ray Core and double down on AIR as an umbrella for our downstream libraries.

It also avoids the question of where to put which modules. Integrations and callbacks can quite naturally remain in ray.air, and both training and tuning modules can access them.

[krfricke]

There's a few more:

• air.callbacks --> also train?
• air.examples for cross-library examples --> docs? (related question: how do we restructure the docs? I guess out of scope for this REP...)
• air.execution --> tentatively train?
• air.util - looks like mostly data?

[pcmoritz]

Thanks for bringing these up :)

air.callbacks has already been moved to air.integrations, which will move to train.integrations, right?

air.examples these should not be living in the source tree, so yes, docs would probably be the right place. I don't consider air.examples to be part of the AIR API -- on the docs: @richardliaw is owning the docs restructuring :)

air.execution these are internal APIs, so are not a concern for this REP. We should feel free to put them into some utility namespace of either train or tune, depending what makes the most sense going forward.

air.util the tensor extension stuff should go into ray.data, the torch stuff into ray.train.torch (I'm a bit confused on whether this is a public API or not -- it seems to be used in the OPT example at the moment? -- we should clarify that and handle it appropriately. Is there anything else here we need to decide about?

[matthewdeng]

Would it make sense to capture all these in the REP to make it more comprehensive, i.e. to describe what it takes to fully "Disband the ray.air namespace"?

[pcmoritz]

On your higher level comment: Not only would introducing ray.air.training, ray.air.tuning (and we would also need ray.air.serving) be a far higher amount of API changes, but it would also lead us to a worse outcome -- everything actually fits neatly into one of the libraries and there is really no need for a ray.air namespace. The libraries ray.train, ray.tune, ray.serve, ray.data and ray.rllib are already well separated from ray core, so I don't think there is danger of confusion here :)

[matthewdeng]

Would it make sense to capture all these in the REP to make it more comprehensive, i.e. to describe what it takes to fully "Disband the ray.air namespace"?

[amogkam]

The only concern is if we are ok with this type of usage for Ray Tune only use case?

from ray import tune
from ray import train

def objective(config):
    score = config["a"] ** 2 + config["b"]
    train.report({"score": score})

search_space = {
    "a": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
    "b": tune.choice([1, 2, 3]),
}

tuner = tune.Tuner(objective, param_space=search_space)

results = tuner.fit()
print(results.get_best_result(metric="score", mode="min").config)

This exposes the ray.train namespace to Tune only users...which may lend to confusion (why do I need to import Ray Train if I am just using Ray Tune?) This was one of the motivations for consolidating under ray.air, so just want to make sure we are considering this case.

[ericl]

Suggested change

	self.predictor = XGBoostPredictor.from_checkpoint(checkpoint)
	local_dir = checkpoint.to_directory()
	self.model: xgboost.Booster = pickle.loads(local_dir + "/saved_booster.pkl")
	``
	Since we are de-emphasizing predictors, shall we also change this to show something like this instead?

[pcmoritz]

I fixed this now by adopting the new way from https://docs.google.com/document/d/1J-09US8cXc-tpl2A1BpOrlHLTEDMdIJp6Ah1ifBUw7Y/edit :)

[krfricke]

Thanks for the updates

[pcmoritz]

@amogkam We decided that it is ok for Tune to depend on Train, since tuning training runs is going to be the most important use case of Ray Tune and also that's how most machine learning engineers / practitioners think about the relation between training and tuning :)

[matthewdeng]

Looks great and feels a lot cleaner!

[ericl]

@zhe-thoughts this can be merged

[zhe-thoughts]

Thanks for the work @pcmoritz @ericl @krfricke @matthewdeng @amogkam