[train] Add fault tolerance variant to the training data ingest benchmark

[justinvyu]

Summary

Adds a .skip_training.fault_tolerance variant to the image classification training ingest release test which kills a node every N seconds and tests the worker recovery.

Also add the ability to stress-test concurrent multi-dataset execution with training ingest (training and validation datasets) by performing validation during the training epoch every N steps. This is not enabled because it is too unperformant and will cause the test to run for too long. This will be addressed in a follow-up PR.

Also updates the training benchmark script to load training state properly and do mid-epoch resumption by skipping batches all the way up to the batch that corresponds to the latest checkpoint.

Adds the following metrics:

• checkpoint/download: Time spent downloading the checkpoint from storage to local.
• checkpoint/load: Time spent loading the checkpoint from local.
• train/iter_skip_batch: Time spent skipping batches upon restoration to do "mid-epoch resumption."
• checkpoint/restoration_time: Extra time spent on restoration, which is just the sum of the above 3.

Here's what the fault tolerance test does:

• Starts a chaos killer which kills a node every ~480 seconds, killing up to 2 nodes across the entire job.
• Runs training with max_failures=4 (with 2 extra failures than needed as a buffer).

[justinvyu]

@raulchen Is this a reasonable way to handle multi-dataset?

[justinvyu]

Note: We should track "restoration time" in Ray Train by default. This time does not include process group re-init, actor startup, etc. Would be good to sum these two together to get the restoration startup cost.

[justinvyu]

Keeping an accurate total for the time spent "skipping" batches is a little tricky because this step could take a long time (ex: if training was killed previously at the second to last batch). Then, a node could get killed in this time, and the time spent skipping batches would be lost from the metrics, which are only snapshotted every checkpoint.

May want to checkpoint the metrics separately from the model if we want to track this better.

[matthewdeng]

Very clean!

[matthewdeng]

nit: use same context manager pattern as done later with the batch iteration?

e.g.

with self._metrics["checkpoint/download"].timer():
    checkpoint.to_directory(temp_checkpoint_dir)


with self._metrics["checkpoint/load"].timer():
    self.load_checkpoint(temp_checkpoint_dir)

[justinvyu]

load_checkpoint will load in the snapshot of the metrics, which will overwrite these guys. So I need to save it separately then add it to the metrics later.

[matthewdeng]

Inverse ordering of timer and condition? Or do you intentionally want to include this as well like you mentioned in the other comment.

            if not self.benchmark_config.skip_validation_step:
                with self._metrics["validation/step"].timer():
                    total_loss += self.validate_step(input_batch, labels)

[justinvyu]

I think it's easier to parse the output if this metric always exists and is just 0 rather than existing conditionally.

[matthewdeng]

nit: Use constants for all keys to make it safe against typos.

[justinvyu]

I'll do this in a followup.