add weights_only arg to checkpoint_io

Author: matsumotosan

src/lightning/fabric/plugins/io/checkpoint_io.py

@@ -47,13 +47,20 @@ def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_optio
         """
 
     @abstractmethod
-    def load_checkpoint(self, path: _PATH, map_location: Optional[Any] = None) -> dict[str, Any]:
+    def load_checkpoint(
+        self, path: _PATH, map_location: Optional[Any] = None, weights_only: bool = True
+    ) -> dict[str, Any]:
         """Load checkpoint from a path when resuming or loading ckpt for test/validate/predict stages.
 
         Args:
             path: Path to checkpoint
             map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage
                 locations.
+            weights_only: If ``True``, restricts loading to ``state_dicts`` of plain ``torch.Tensor`` and other
+            primitive types. If loading a checkpoint from a trusted source that contains an ``nn.Module``, use
+                ``weights_only=False``. If loading checkpoint from an untrusted source, we recommend using
+                ``weights_only=True``. For more information, please refer to the
+                `PyTorch Developer Notes on Serialization Semantics <https://docs.pytorch.org/docs/main/notes/serialization.html#id3>`__.
 
         Returns: The loaded checkpoint.
 

src/lightning/fabric/plugins/io/torch_io.py

@@ -59,7 +59,7 @@ def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_optio
 
     @override
     def load_checkpoint(
-        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage
+        self, path: _PATH, map_location: Optional[Callable] = lambda storage, loc: storage, weights_only: bool = True
     ) -> dict[str, Any]:
         """Loads checkpoint using :func:`torch.load`, with additional handling for ``fsspec`` remote loading of files.
 
@@ -80,7 +80,7 @@ def load_checkpoint(
         if not fs.exists(path):
             raise FileNotFoundError(f"Checkpoint file not found: {path}")
 
-        return pl_load(path, map_location=map_location)
+        return pl_load(path, map_location=map_location, weights_only=weights_only)
 
     @override
     def remove_checkpoint(self, path: _PATH) -> None:

tests/legacy/generate_checkpoints.sh

@@ -16,9 +16,9 @@ printf "PYTHONPATH: $PYTHONPATH"
 rm -rf $ENV_PATH
 
 function create_and_save_checkpoint {
-  python --version
-  python -m pip --version
-  python -m pip list
+  # python --version
+  # python -m pip --version
+  # python -m pip list
 
   python $LEGACY_FOLDER/simple_classif_training.py $pl_ver
 
@@ -52,10 +52,12 @@ done
 if [[ -z "$@" ]]; then
   printf "\n\n processing local version\n"
 
-  python -m pip install \
+  # python -m pip install \
+  uv pip install \
     -r $LEGACY_FOLDER/requirements.txt \
     -r "$(dirname $TESTS_FOLDER)/requirements/pytorch/test.txt" \
     -f https://download.pytorch.org/whl/cpu/torch_stable.html
   pl_ver="local"
+  # pl_ver=$(python -c "import lightning.pytorch as pl; print(pl.__version__)")
   create_and_save_checkpoint
 fi

tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py

@@ -32,7 +32,9 @@ class CustomCheckpointIO(CheckpointIO):
     def save_checkpoint(self, checkpoint: dict[str, Any], path: _PATH, storage_options: Optional[Any] = None) -> None:
         torch.save(checkpoint, path)
 
-    def load_checkpoint(self, path: _PATH, storage_options: Optional[Any] = None) -> dict[str, Any]:
+    def load_checkpoint(
+        self, path: _PATH, storage_options: Optional[Any] = None, weights_only: bool = True
+    ) -> dict[str, Any]:
         return torch.load(path, weights_only=True)
 
     def remove_checkpoint(self, path: _PATH) -> None: