[AIR] <Part 2> Support metric logging and checkpointing for LightningTrainer

[woshiyyya]

Why are these changes needed?

There will be a list of PRs for PyTorch Lightning Integration. This is the second one.

• LightningTrainer + Test
• Logging and Checkpointing + Test
• Predictor + Batch Prediction + Test
• Integrate with Tune + Test

Content for this PR:

• Implemented RayModelCheckpoint callback to report latest metrics and checkpoint to trainer session
• Implemented LightningCheckpoint, a subclass of TorchCheckpoint.
• CI Tests

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

[justinvyu]

Nice! I have a few questions:

Also, do we need to implement a custom restore method, or is using DataParallelTrainer's default one ok? The only PTL-specific that you may need to re-specify is lightning_config. I.e. re-specifying the lightning module class. Can talk offline about this.

[gjoliver]

pretty nice. a few questions.

[gjoliver]

why is metrics reporting part of the checkpoint class?
what if I want to report data / iteration, but don't want to create checkpoints?

[woshiyyya]

The context here is, checkpointing and logging are separate logics in Lightning. Checkpoint class can access metrics and checkpoint, but Logger can only access metrics. In order to report checkpoint and metrics together, we implement reporting in checkpoint class.

For logging, we recommend the users keep using lightning's native Loggers(e.g. wandb, mlflow, tensorboard loggers). They can control the logging frequency by themselves and retrieve logs as usual, which is less intrusive and aligns better with user habits.

[gjoliver]

don't feel particularly safe about this. why do we have this keyword, and it's not even __ prefixed ...
also we may be logging this warning msg every time this is called.
is there a list of such keywords defined somewhere?

[woshiyyya]

I agree. I tried but didn't relevant keys in ray.air.constants. I removed this warning msg and changed the key name to __report_on now.

[gjoliver]

wait, do we want people to pass in a dict or the raw Builder()?
the type annotation says it's a raw dict already.

[woshiyyya]

We expect them to pass a dict, but they can also choose not to and specify it later in Tuner's param_space. We will check if lightning_config is provided in the training_loop_per_worker.

For example:

lightning_config = (
    LightningConfigBuilder()
    .module(MNISTClassifier, feature_dim=128, lr=tune.grid_search([0.1, 0.01, 0.001]))
    .trainer(max_epochs=3, accelerator="cpu")
    .fit_params(datamodule=datamodule)
    .build()
)

trainer = LightningTrainer(scaling_config=scaling_config)

tuner = tune.Tuner(
    trainer,
    param_space={"lightning_config": lightning_config}
    ...
)

[gjoliver]

😓
any chance we can disable the csv writer so we don't have to do this hack here?

[woshiyyya]

Oh actually this metrics check is not in csv writer, but in trial_runner and is triggered every time we call session.report. I agree this is hacky, but it seems to be the only way we can skip this check now. 😅

[richardliaw]

stamp for docs

[ollie-iterators]

This code is causing documentation to fail