Feat: checkpoint optimize function to restart after crash

[deependujha]

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes #137 .

Make optimize function fault-tolerant and restart from last checkpoint after an unexpected crash.

The checkpoint feature is supported for mode = "None | append". Overwrite mode will overwrite everything, even checkpoints.

If you were executing Optimize function in overwrite mode with use_checkpoint=True, and if it crashes.
Execute with use_checkpoint=True & mode=None in the next run.

• Sample code for this feature

from litdata import optimize, StreamingDataset

def fn(i: int):
    if i in [22, 36, 59, 78]:
        raise ValueError("An error occurred")

    return i, i**2

def another_fn(i: int):
    return i, i**2

output_dir = "s3://my-dummy-bucket-litdata/is-litdata-checkpoint-working/"
output_dir = "data"

try:
    optimize(
        fn=fn,
        inputs=list(range(100)),
        output_dir=output_dir,
        chunk_size=5,
        num_workers=4,
        use_checkpoint=True,
    )
except Exception as _:
    print("fn failed. Now, we will go for another_fn.")
    optimize(
        fn=another_fn,
        inputs=list(range(100)),
        output_dir=output_dir,
        chunk_size=5,
        num_workers=4,
        use_checkpoint=True,
    )

ds = StreamingDataset(output_dir)
print(f"{len(ds)=}")
print(f"{ds[:]=}")
print("All done ✅")

---

🚧 Warning: It doesn't support generator functions as of now. It'll be added in the next PR.

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Attention: Patch coverage is 53.57143% with 91 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@1616aeb). Learn more about missing BASE report.

Additional details and impacted files

@@          Coverage Diff          @@
##             main   #206   +/-   ##
=====================================
  Coverage        ?    77%           
=====================================
  Files           ?     33           
  Lines           ?   4667           
  Branches        ?      0           
=====================================
  Hits            ?   3588           
  Misses          ?   1079           
  Partials        ?      0           

[tchaton]

Looks great !