Feature/handle pytorch load changes (#306)

* Introduce default parameter weights_only to load_checkpoint to be aligned with changes in PyTorch Lightning
* Align test types with torch.load weights_only support

Author: IsaevIlya

s3torchconnector/src/s3torchconnector/lightning/s3_lightning_checkpoint.py

@@ -73,7 +73,12 @@ def load_checkpoint(
         s3reader = self._client.get_object(bucket, key)
         # FIXME - io.BufferedIOBase and typing.IO aren't compatible
         #  See https://github.com/python/typeshed/issues/6077
-        return torch.load(s3reader, map_location)  # type: ignore
+
+        # Explicitly set weights_only=False to:
+        # 1. Maintain backwards compatibility with older PyTorch versions where this was the default behavior
+        # 2. Match PyTorch Lightning's implementation strategy for consistent behavior
+        # Reference: https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/utilities/cloud_io.py#L36
+        return torch.load(s3reader, map_location, weights_only=False)  # type: ignore
 
     def remove_checkpoint(
         self,

s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py

@@ -57,7 +57,7 @@ def test_save_compatibility_with_s3_checkpoint(checkpoint_directory):
     s3_uri = f"{checkpoint_directory.s3_uri}{checkpoint_name}"
     s3_lightning_checkpoint.save_checkpoint(tensor, s3_uri)
     checkpoint = S3Checkpoint(region=checkpoint_directory.region)
-    loaded_checkpoint = torch.load(checkpoint.reader(s3_uri))
+    loaded_checkpoint = torch.load(checkpoint.reader(s3_uri), weights_only=False)
     assert torch.equal(tensor, loaded_checkpoint)
 
 

s3torchconnector/tst/e2e/test_e2e_s3checkpoint.py

@@ -20,7 +20,7 @@ def test_general_checkpointing(checkpoint_directory, tensor_dimensions):
     with checkpoint.writer(s3_uri) as writer:
         torch.save(tensor, writer)
 
-    loaded = torch.load(checkpoint.reader(s3_uri))
+    loaded = torch.load(checkpoint.reader(s3_uri), weights_only=True)
 
     assert torch.equal(tensor, loaded)
 
@@ -49,7 +49,7 @@ def test_nn_checkpointing(checkpoint_directory):
     # assert models are not equal before loading from checkpoint
     assert not nn_model.equals(loaded_nn_model)
 
-    loaded_checkpoint = torch.load(checkpoint.reader(s3_uri))
+    loaded_checkpoint = torch.load(checkpoint.reader(s3_uri), weights_only=True)
     loaded_nn_model.load_state_dict(loaded_checkpoint["model_state_dict"])
     assert nn_model.equals(loaded_nn_model)
 

s3torchconnector/tst/unit/_checkpoint_byteorder_patch.py

@@ -23,4 +23,4 @@ def save_with_byteorder(data, fobj, byteorder: str, use_modern_pytorch_format: b
 
 def load_with_byteorder(fobj, byteorder):
     with _patch_byteorder(byteorder):
-        return torch.load(fobj)
+        return torch.load(fobj, weights_only=True)

s3torchconnector/tst/unit/_hypothesis_python_primitives.py

@@ -3,47 +3,30 @@
 
 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()
+    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))
-)
+hashable = deferred(lambda: (scalars | 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)
     )