[Train] LightningTrainer failed to auto-restore if ckpt_path provided in LightningConfigbuilder.fit_param()

[woshiyyya]

What happened + What you expected to happen

If users provide a ckpt_path in LightningConfigBuilder.fit_param(), when resuming training from a failed run, LightningTrainer will always resume training from that ckpt_path again, instead of the latest AIR checkpoint.

Code:

ray/python/ray/train/lightning/lightning_trainer.py

Lines 543 to 549 in 0fd06ad

	checkpoint = session.get_checkpoint()
	if checkpoint and "ckpt_path" not in trainer_fit_params:
	with checkpoint.as_directory() as ckpt_dir:
	trainer_fit_params["ckpt_path"] = f"{ckpt_dir}/{MODEL_KEY}"
	trainer.fit(lightning_module, **trainer_fit_params)
	else:
	trainer.fit(lightning_module, **trainer_fit_params)

Instead, we should always restore from the latest AIR checkpoint, if we can get it from the session, regardless of whether the user provided a ckpt_path or not.

Versions / Dependencies

master

Reproduction script

lightning_config = (
    LightningConfigBuilder()
    .module(
        ...
    )
    .trainer(
        ...
    )
    .fit_params(
        train_dataloaders=train_loader,
        val_dataloaders=val_loader,
        ckpt_path="YOUR_CKPT_PATH"
    )
    .build()
)

Together with a FailureConfig

run_config = RunConfig(
    ...,
    failure_config=FailureConfig(max_failures=3), 
)

Issue Severity

High: It blocks me from completing my task.