[train] expose training input/output in callbacks

[matthewdeng]

Changes

• Made TrainRunContext frozen and added all training inputs (configs, datasets, etc.)
• Changed TrainContext to contain TrainRunContext instead of inheriting from it
• Added after_controller_finish callback to expose final training results
• Modified after_controller_start to receive TrainRunContext parameter

TODO

• [ ] Remove redundant TrainRunContext from Callback initializations
• [ ] Refactor logging to handle both TrainContext and TrainRunContext
• [ ] Fix/add tests