Added support for passing a custom endpoint to S3LightningCheckpoint (#201)

* Added support for passing a custom endpoint to S3LightningCheckpoint
* Updated CHANGELOG

Author: bharatendra

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### New features
 * Update S3ClientConfig to pass in the configuration for allowing unsigned requests, under boolean flag `unsigned`.
 * Improve the performance of `s3reader` when utilized with `pytorch.load` by incorporating support for the `readinto` method.
+* Add support for passing an optional custom endpoint to S3LightningCheckpoint constructor method.
 
 
 ## v1.2.2 (March 22, 2024)

s3torchconnector/src/s3torchconnector/lightning/s3_lightning_checkpoint.py

@@ -20,13 +20,15 @@ def __init__(
         self,
         region: str,
         s3client_config: Optional[S3ClientConfig] = None,
+        endpoint: Optional[str] = None,
     ):
         self.region = region
         user_agent = UserAgent(["lightning", lightning.__version__])
         self._client = S3Client(
             region,
             user_agent=user_agent,
             s3client_config=s3client_config,
+            endpoint=endpoint,
         )
 
     def save_checkpoint(

s3torchconnector/tst/unit/lightning/test_s3_lightning_checkpoint.py

@@ -27,6 +27,7 @@
 TEST_BUCKET = "test-bucket"
 TEST_KEY = "test-key"
 TEST_REGION = "us-east-1"
+TEST_ENDPOINT = "https://s3.us-east-1.amazonaws.com"
 
 
 @pytest.fixture()
@@ -148,6 +149,11 @@ def test_teardown(lightning_checkpoint):
     # Assert no exception is thrown - implicit
 
 
+def test_lightning_checkpoint_creation_with_region_and_endpoint():
+    checkpoint = S3LightningCheckpoint(TEST_REGION, endpoint=TEST_ENDPOINT)
+    assert isinstance(checkpoint, S3LightningCheckpoint)
+
+
 def _test_save(
     client,
     checkpoint: CheckpointIO,