Add examples for reading/writing checkpoints with lightning in various ways (#150)

* Add examples for reading/writing checkpoints with lightning in various ways

* Add small set of documentation for lightning support to the README

---------

Co-authored-by: Simon Beal <[email protected]>

Author: muddyfish

README.md

@@ -115,6 +115,40 @@ For example, assuming the following directory bucket name `my-test-bucket--usw2-
 usw2-az1, then the URI used will look like: `s3://my-test-bucket--usw2-az1--x-s3/<PREFIX>` (**please note that the 
 prefix for Amazon S3 Express One Zone should end with '/'**), paired with region us-west-2.
 
+
+## Lightning Integration
+
+Amazon S3 Connector for PyTorch includes an integration for PyTorch Lightning, featuring S3LightningCheckpoint, an 
+implementation of Lightning's CheckpointIO. This allows users to make use of Amazon S3 Connector for PyTorch's S3 
+checkpointing functionality with Pytorch Lightning.
+
+### Getting Started
+
+#### Installation
+
+```sh
+pip install s3torchconnector[lightning]
+```
+
+### Examples
+
+End to end examples for the Pytorch Lightning integration can be found in the 
+[examples/lightning](https://github.com/awslabs/s3-connector-for-pytorch/tree/main/examples/lightning) directory
+
+```py
+from lightning import Trainer
+from s3torchconnector.lightning import S3LightningCheckpoint
+
+...
+
+s3_checkpoint_io = S3LightningCheckpoint("us-east-1")
+trainer = Trainer(
+    plugins=[s3_checkpoint_io],
+    default_root_dir="s3://bucket_name/key_prefix/"
+)
+trainer.fit(model)
+```
+
 ## Contributing
 We welcome contributions to Amazon S3 Connector for PyTorch. Please 
 see [CONTRIBUTING](https://github.com/awslabs/s3-connector-for-pytorch/blob/main/doc/CONTRIBUTING.md) 

examples/lightning/async_checkpoint_writing.py

@@ -0,0 +1,44 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from lightning import Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.demos import WikiText2, LightningTransformer
+from lightning.pytorch.plugins import AsyncCheckpointIO
+from torch.utils.data import DataLoader
+
+from s3torchconnector.lightning import S3LightningCheckpoint
+
+
+def main(region: str, checkpoint_path: str):
+    dataset = WikiText2()
+    dataloader = DataLoader(dataset, num_workers=3)
+
+    model = LightningTransformer(vocab_size=dataset.vocab_size)
+    s3_lightning_checkpoint = S3LightningCheckpoint(region)
+    async_checkpoint = AsyncCheckpointIO(s3_lightning_checkpoint)
+
+    # This will create one checkpoint per 'step', which we define later to be 8.
+    # To checkpoint more or less often, change `every_n_train_steps`.
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=checkpoint_path,
+        save_top_k=-1,
+        every_n_train_steps=1,
+        filename="checkpoint-{epoch:02d}-{step:02d}",
+        enable_version_counter=True,
+    )
+
+    trainer = Trainer(
+        plugins=[async_checkpoint],
+        callbacks=[checkpoint_callback],
+        min_epochs=4,
+        max_epochs=8,
+        max_steps=8,
+    )
+    trainer.fit(model, dataloader)
+
+
+if __name__ == "__main__":
+    import os
+
+    main(os.getenv("REGION"), os.getenv("CHECKPOINT_PATH"))

examples/lightning/checkpoint_manual_save.py

@@ -0,0 +1,34 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from lightning import Trainer
+from lightning.pytorch.demos import WikiText2, LightningTransformer
+from torch.utils.data import DataLoader
+
+from s3torchconnector.lightning import S3LightningCheckpoint
+
+
+def main(region: str, checkpoint_path: str):
+    dataset = WikiText2()
+    dataloader = DataLoader(dataset, num_workers=3)
+
+    model = LightningTransformer(vocab_size=dataset.vocab_size)
+    s3_lightning_checkpoint = S3LightningCheckpoint(region)
+
+    # No automatic checkpointing set up here.
+    trainer = Trainer(
+        plugins=[s3_lightning_checkpoint],
+        enable_checkpointing=False,
+        min_epochs=4,
+        max_epochs=5,
+        max_steps=3,
+    )
+    trainer.fit(model, dataloader)
+    # Manually create checkpoint to the desired location
+    trainer.save_checkpoint(checkpoint_path)
+
+
+if __name__ == "__main__":
+    import os
+
+    main(os.getenv("REGION"), os.getenv("CHECKPOINT_PATH"))

examples/lightning/checkpoint_reading.py

@@ -0,0 +1,31 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from lightning import Trainer
+from lightning.pytorch.demos import WikiText2, LightningTransformer
+from torch.utils.data import DataLoader
+
+from s3torchconnector.lightning import S3LightningCheckpoint
+
+
+def main(region: str, checkpoint_path: str):
+    dataset = WikiText2()
+    dataloader = DataLoader(dataset, num_workers=3)
+
+    model = LightningTransformer(vocab_size=dataset.vocab_size)
+    s3_lightning_checkpoint = S3LightningCheckpoint(region)
+
+    trainer = Trainer(
+        plugins=[s3_lightning_checkpoint],
+        min_epochs=4,
+        max_epochs=5,
+        max_steps=3,
+    )
+    # Load the checkpoint in `ckpt_path` before training
+    trainer.fit(model, dataloader, ckpt_path=checkpoint_path)
+
+
+if __name__ == "__main__":
+    import os
+
+    main(os.getenv("REGION"), os.getenv("CHECKPOINT_PATH"))

examples/lightning/checkpoint_writing.py

@@ -0,0 +1,47 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from lightning import Trainer
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.demos import WikiText2, LightningTransformer
+from torch.utils.data import DataLoader
+
+from s3torchconnector.lightning import S3LightningCheckpoint
+
+
+def main(region: str, checkpoint_path: str, save_only_latest: bool):
+    dataset = WikiText2()
+    dataloader = DataLoader(dataset, num_workers=3)
+
+    model = LightningTransformer(vocab_size=dataset.vocab_size)
+    s3_lightning_checkpoint = S3LightningCheckpoint(region)
+
+    # Save once per step, and if `save_only_latest`, replace the last checkpoint each time.
+    # Replacing is implemented by saving the new checkpoint, and then deleting the previous one.
+    # If `save_only_latest` is False, a new checkpoint is created for each step.
+    checkpoint_callback = ModelCheckpoint(
+        dirpath=checkpoint_path,
+        save_top_k=1 if save_only_latest else -1,
+        every_n_train_steps=1,
+        filename="checkpoint-{epoch:02d}-{step:02d}",
+        enable_version_counter=True,
+    )
+
+    trainer = Trainer(
+        plugins=[s3_lightning_checkpoint],
+        callbacks=[checkpoint_callback],
+        min_epochs=4,
+        max_epochs=5,
+        max_steps=3,
+    )
+    trainer.fit(model, dataloader)
+
+
+if __name__ == "__main__":
+    import os
+
+    main(
+        os.getenv("REGION"),
+        os.getenv("CHECKPOINT_PATH"),
+        os.getenv("LATEST_CHECKPOINT_ONLY") == "1",
+    )