Lightning support (#132)

- Add testing which splits out utils for checkpointing into a common module
- Add lightning optional dependency list for lightning checkpoints
- Add delete tests
- Add tests that validate the error message we get if lightning is not installed

---------

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

Author: 2 people

.github/workflows/python-checks.yml

@@ -60,6 +60,12 @@ jobs:
       - name: s3torchbenchmarking unit tests
         run: pytest s3torchbenchmarking/tst --hypothesis-profile ci --hypothesis-show-statistics -c ./
 
+      - name: Install Lightning dependency
+        run: |
+          python -m pip install -e "s3torchconnector[lightning]"
+      - name: s3torchconnector lightning unit tests
+        run: pytest s3torchconnector/tst/unit/lightning --hypothesis-profile ci --hypothesis-show-statistics
+
   lint:
     name: Python lints
     runs-on: ubuntu-22.04

s3torchconnector/pyproject.toml

@@ -45,6 +45,10 @@ e2e = [
     "pytest-xdist"
 ]
 
+lightning = [
+    "lightning >= 2.0"
+]
+
 [tool.setuptools.packages]
 # Pure Python packages/modules
 find = { where = ["src"] }

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -105,3 +105,7 @@ def list_objects(
     def head_object(self, bucket: str, key: str) -> ObjectInfo:
         log.debug(f"HeadObject s3://{bucket}/{key}")
         return self._client.head_object(bucket, key)
+
+    def delete_object(self, bucket: str, key: str) -> None:
+        log.debug(f"DeleteObject s3://{bucket}/{key}")
+        self._client.delete_object(bucket, key)

s3torchconnector/src/s3torchconnector/lightning/__init__.py

@@ -0,0 +1,11 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+#  Get a nice error message if lightning isn't available.
+import lightning
+
+from .s3_lightning_checkpoint import S3LightningCheckpoint
+
+__all__ = [
+    "S3LightningCheckpoint",
+]

s3torchconnector/src/s3torchconnector/lightning/s3_lightning_checkpoint.py

@@ -0,0 +1,83 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from typing import Optional, Dict, Any
+
+import torch
+
+from lightning.pytorch.plugins.io import CheckpointIO
+
+from .._s3client import S3Client
+from .._s3dataset_common import parse_s3_uri
+
+
+class S3LightningCheckpoint(CheckpointIO):
+    """A checkpoint manager for S3 using the :class:`CheckpointIO` interface."""
+
+    def __init__(self, region: str):
+        self.region = region
+        self._client = S3Client(region)
+
+    def save_checkpoint(
+        self,
+        checkpoint: Dict[str, Any],
+        path: str,
+        storage_options: Optional[Any] = None,
+    ) -> None:
+        """Save model/training states as a checkpoint file through state-dump and upload to S3.
+
+        Args:
+            checkpoint (Dict[str, Any]): Containing model and trainer state
+            path (str): Write-target S3 uri
+            storage_options: Optional parameters when saving the model/training states.
+        """
+        self._validate_path(path)
+        bucket, key = parse_s3_uri(path)
+        with self._client.put_object(bucket, key) as s3writer:
+            torch.save(checkpoint, s3writer)
+
+    def load_checkpoint(
+        self,
+        path: str,
+        map_location: Optional[Any] = None,
+    ) -> Dict[str, Any]:
+        """Load checkpoint from an S3 location when resuming or loading ckpt for test/validate/predict stages.
+
+        Args:
+            path (str): S3 uri to checkpoint
+            map_location: A function, :class:`torch.device`, string or a dict specifying how to remap storage locations.
+
+        Returns:
+            Dict[str, Any]: The loaded checkpoint
+
+        Raises:
+            S3Exception: An error occurred accessing S3.
+        """
+        self._validate_path(path)
+        bucket, key = parse_s3_uri(path)
+        s3reader = self._client.get_object(bucket, key)
+        return torch.load(s3reader, map_location)
+
+    def remove_checkpoint(self, path: str) -> None:
+        """Remove checkpoint file from the S3 uri.
+
+        Args:
+            path (str): S3 uri to checkpoint
+
+        Raises:
+            S3Exception: An error occurred accessing S3.
+        """
+        self._validate_path(path)
+        bucket, key = parse_s3_uri(path)
+        self._client.delete_object(bucket, key)
+
+    def teardown(self) -> None:
+        """This method is called to teardown the process."""
+        pass
+
+    @staticmethod
+    def _validate_path(path: str) -> None:
+        if not isinstance(path, str):
+            raise TypeError(
+                f"{type(path).__name__!r} is not a supported type for 'path'. Must be a string formatted as an S3 uri."
+            )

s3torchconnector/tst/unit/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD

s3torchconnector/tst/unit/_checkpoint_byteorder_patch.py

@@ -0,0 +1,26 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from contextlib import contextmanager
+from unittest.mock import patch
+import torch
+from hypothesis.strategies import one_of, just
+
+byteorders = one_of(just("little"), just("big"))
+
+
+@contextmanager
+def _patch_byteorder(byteorder: str):
+    with patch("torch.serialization.sys") as mock_sys:
+        mock_sys.byteorder = byteorder
+        yield
+
+
+def save_with_byteorder(data, fobj, byteorder: str, use_modern_pytorch_format: bool):
+    with _patch_byteorder(byteorder):
+        torch.save(data, fobj, _use_new_zipfile_serialization=use_modern_pytorch_format)
+
+
+def load_with_byteorder(fobj, byteorder):
+    with _patch_byteorder(byteorder):
+        return torch.load(fobj)

s3torchconnector/tst/unit/_hypothesis_python_primitives.py

@@ -0,0 +1,50 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from hypothesis.strategies import (
+    integers,
+    binary,
+    none,
+    characters,
+    complex_numbers,
+    floats,
+    booleans,
+    decimals,
+    fractions,
+    deferred,
+    frozensets,
+    tuples,
+    dictionaries,
+    lists,
+    uuids,
+    sets,
+    text,
+)
+
+scalars = (
+    none()
+    | booleans()
+    | integers()
+    # Disallow nan as it doesn't have self-equality
+    | floats(allow_nan=False)
+    | complex_numbers(allow_nan=False)
+    | decimals(allow_nan=False)
+    | fractions()
+    | characters()
+    | binary(max_size=10)
+    | text(max_size=10)
+    | uuids()
+)
+
+hashable = deferred(
+    lambda: (scalars | frozensets(hashable, max_size=5) | tuples(hashable))
+)
+
+python_primitives = deferred(
+    lambda: (
+        hashable
+        | sets(hashable, max_size=5)
+        | lists(python_primitives, max_size=5)
+        | dictionaries(keys=hashable, values=python_primitives, max_size=3)
+    )
+)

s3torchconnector/tst/unit/lightning/__init__.py

@@ -0,0 +1,2 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD

s3torchconnector/tst/unit/lightning/test_s3_lightning_checkpoint.py

@@ -0,0 +1,185 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+from io import BytesIO
+from operator import eq
+from pathlib import Path
+from typing import Callable, Any
+
+import hypothesis
+import pytest
+import torch
+from hypothesis import given, HealthCheck
+from lightning.fabric.plugins import CheckpointIO
+from lightning.pytorch.plugins import AsyncCheckpointIO
+
+from s3torchconnector._s3client import MockS3Client
+from s3torchconnector.lightning import S3LightningCheckpoint
+from .._checkpoint_byteorder_patch import (
+    byteorders,
+    save_with_byteorder,
+    load_with_byteorder,
+    _patch_byteorder,
+)
+from .._hypothesis_python_primitives import python_primitives
+from s3torchconnectorclient import S3Exception
+
+TEST_BUCKET = "test-bucket"
+TEST_KEY = "test-key"
+TEST_REGION = "us-east-1"
+
+
+@pytest.fixture()
+def client() -> MockS3Client:
+    return MockS3Client(TEST_REGION, TEST_BUCKET)
+
+
+@pytest.fixture()
+def lightning_checkpoint(client) -> S3LightningCheckpoint:
+    s3_lightning_checkpoint = S3LightningCheckpoint(TEST_REGION)
+    s3_lightning_checkpoint._client = client
+    return s3_lightning_checkpoint
+
+
+@pytest.fixture()
+def async_lightning_checkpoint(lightning_checkpoint) -> Callable[[], AsyncCheckpointIO]:
+    return lambda: AsyncCheckpointIO(lightning_checkpoint)
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(python_primitives, byteorders)
+def test_lightning_checkpointing_saves_python_primitives(
+    client, lightning_checkpoint, data, byteorder
+):
+    _test_save(client, lightning_checkpoint, data, byteorder)
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(byteorders)
+def test_lightning_checkpointing_saves_tensor(client, lightning_checkpoint, byteorder):
+    tensor = torch.rand(2, 4)
+    _test_save(client, lightning_checkpoint, tensor, byteorder, equal=torch.equal)
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(byteorders)
+def test_async_lightning_checkpointing_saves_tensor(
+    client, async_lightning_checkpoint, byteorder
+):
+    tensor = torch.rand(2, 4)
+    _test_save(
+        client, async_lightning_checkpoint(), tensor, byteorder, equal=torch.equal
+    )
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(byteorders)
+def test_lightning_checkpointing_saves_untyped_storage(
+    client, lightning_checkpoint, byteorder
+):
+    storage = torch.UntypedStorage([1, 2, 3])
+    _test_save(
+        client,
+        lightning_checkpoint,
+        storage,
+        byteorder,
+        equal=lambda a, b: list(a) == list(b),
+    )
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(python_primitives, byteorders)
+def test_lightning_checkpointing_loads_python_primitives(
+    client, lightning_checkpoint, data, byteorder
+):
+    _test_load(client, lightning_checkpoint, data, byteorder)
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(byteorders)
+def test_lightning_checkpointing_loads_tensor(client, lightning_checkpoint, byteorder):
+    tensor = torch.rand(2, 4)
+    _test_load(client, lightning_checkpoint, tensor, byteorder, equal=torch.equal)
+
+
+@hypothesis.settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
+@given(byteorders)
+def test_lightning_checkpointing_loads_untyped_storage(
+    client, lightning_checkpoint, byteorder
+):
+    storage = torch.UntypedStorage([1, 2, 3])
+    _test_load(
+        client,
+        lightning_checkpoint,
+        storage,
+        byteorder,
+        equal=lambda a, b: list(a) == list(b),
+    )
+
+
+def test_removes_checkpoint(client, lightning_checkpoint):
+    lightning_checkpoint.remove_checkpoint(f"s3://{TEST_BUCKET}/{TEST_KEY}")
+
+    with pytest.raises(S3Exception) as error:
+        client.get_object(TEST_BUCKET, TEST_KEY).read()
+    assert str(error.value) == "Service error: The key does not exist"
+
+
+@pytest.mark.parametrize(
+    "checkpoint_method_name, kwargs",
+    [
+        ("save_checkpoint", {"path": Path("bucket", "key"), "checkpoint": None}),
+        ("load_checkpoint", {"path": Path("/", "bucket", "key")}),
+        ("remove_checkpoint", {"path": Path()}),
+        ("remove_checkpoint", {"path": ["not", "a", "string"]}),
+    ],
+)
+def test_invalid_path(lightning_checkpoint, checkpoint_method_name, kwargs):
+    checkpoint_method = getattr(lightning_checkpoint, checkpoint_method_name)
+    with pytest.raises(
+        TypeError,
+        match="is not a supported type for 'path'. Must be a string formatted as an S3 uri",
+    ):
+        checkpoint_method(**kwargs)
+
+
+def test_teardown(lightning_checkpoint):
+    lightning_checkpoint.teardown()
+    # Assert no exception is thrown - implicit
+
+
+def _test_save(
+    client,
+    checkpoint: CheckpointIO,
+    data,
+    byteorder: str,
+    *,
+    equal: Callable[[Any, Any], bool] = eq,
+):
+    with _patch_byteorder(byteorder):
+        checkpoint.save_checkpoint(data, f"s3://{TEST_BUCKET}/{TEST_KEY}")
+        # For async checkpointing, ensure that we finish writing the checkpoint before we un-patch the byteorder
+        checkpoint.teardown()
+
+    serialised = BytesIO(b"".join(client.get_object(TEST_BUCKET, TEST_KEY)))
+    assert equal(load_with_byteorder(serialised, byteorder), data)
+
+
+def _test_load(
+    client,
+    checkpoint: CheckpointIO,
+    data,
+    byteorder: str,
+    *,
+    equal: Callable[[Any, Any], bool] = eq,
+):
+    # Put some data to mock bucket and use mock client
+    serialised = BytesIO()
+    save_with_byteorder(data, serialised, byteorder, use_modern_pytorch_format=True)
+    serialised.seek(0)
+    client.add_object(TEST_KEY, serialised.read())
+
+    with _patch_byteorder(byteorder):
+        returned_data = checkpoint.load_checkpoint(f"s3://{TEST_BUCKET}/{TEST_KEY}")
+
+    assert equal(returned_data, data)