Merge pull request #394 from VectorInstitute/ft/state-checkpointing-refactor

State checkpointer refactor

Author: fatemetkl

examples/nnunet_example/client.py

@@ -5,8 +5,8 @@
 from os.path import exists, join
 from pathlib import Path
 
-from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer
 from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule
+from fl4health.checkpointing.state_checkpointer import ClientStateCheckpointer
 
 with warnings.catch_warnings():
     # Silence deprecation warnings from sentry sdk due to flwr and wandb
@@ -73,7 +73,7 @@ def main(
 
     if intermediate_client_state_dir is not None:
         checkpoint_and_state_module = ClientCheckpointAndStateModule(
-            state_checkpointer=PerRoundStateCheckpointer(Path(intermediate_client_state_dir))
+            state_checkpointer=ClientStateCheckpointer(Path(intermediate_client_state_dir))
         )
     else:
         checkpoint_and_state_module = None
@@ -147,7 +147,7 @@ def main(
             even if the preprocessed data is found to already exist",
     )
     parser.add_argument(
-        "--server_address",
+        "--server-address",
         type=str,
         required=False,
         default="0.0.0.0:8080",

examples/nnunet_example/server.py

@@ -12,8 +12,8 @@
 from flwr.server.client_manager import SimpleClientManager
 from flwr.server.strategy import FedAvg
 
-from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer
 from fl4health.checkpointing.server_module import NnUnetServerCheckpointAndStateModule
+from fl4health.checkpointing.state_checkpointer import NnUnetServerStateCheckpointer
 from fl4health.metrics.metric_aggregation import evaluate_metrics_aggregation_fn, fit_metrics_aggregation_fn
 from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
 from fl4health.servers.nnunet_server import NnunetServer
@@ -85,7 +85,7 @@ def main(
     )
 
     state_checkpointer = (
-        PerRoundStateCheckpointer(Path(intermediate_server_state_dir))
+        NnUnetServerStateCheckpointer(Path(intermediate_server_state_dir))
         if intermediate_server_state_dir is not None
         else None
     )
@@ -132,7 +132,7 @@ def main(
     )
 
     parser.add_argument(
-        "--intermediate-server-state-dir",
+        "--intermediate-server-state_dir",
         type=str,
         required=False,
         default=None,

fl4health/checkpointing/checkpointer.py

@@ -2,8 +2,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from logging import ERROR, INFO, WARNING
-from pathlib import Path
-from typing import Any
 
 import torch
 import torch.nn as nn
@@ -276,7 +274,6 @@ def __init__(
         maximize: bool = False,
     ) -> None:
         """Checkpointer that checkpoints based on the value of a user defined metric.
-
         Args:
             checkpoint_dir (str): Directory to which the model is saved. This directory should already exist. The
                 checkpointer will not create it if it does not.
@@ -310,83 +307,3 @@ def metric_score_function(_: float, metrics: dict[str, Scalar]) -> float:
             checkpoint_score_name=metric,
             maximize=maximize,
         )
-
-
-class PerRoundStateCheckpointer:
-    def __init__(self, checkpoint_dir: Path) -> None:
-        """
-        Base class that provides a uniform interface for loading, saving and checking if checkpoints exists.
-
-        Args:
-            checkpoint_dir (Path): Base directory to store checkpoints. This checkpoint directory MUST already exist.
-                It will not be created by this state checkpointer.
-        """
-        log(
-            WARNING,
-            "Creating PerRoundCheckpointer. Currently, this functionality is still experimental and only supported "
-            "for BasicClient and NnunetClient, along with their associated servers.",
-        )
-        self.checkpoint_dir = checkpoint_dir
-
-    def save_checkpoint(self, checkpoint_name: str, checkpoint_dict: dict[str, Any]) -> None:
-        """
-        Saves ``checkpoint_dict`` to checkpoint path form from this classes checkpointer dir and the provided
-        checkpoint name.
-
-        Args:
-            checkpoint_name (str): Name of the state checkpoint file.
-            checkpoint_dict (dict[str, Any]): A dictionary with string keys and values of type Any representing the
-                state to checkpoint.
-
-        Raises:
-            e: Will throw an error if there is an issue saving the model. ``Torch.save`` seems to swallow errors in
-                this context, so we explicitly surface the error with a try/except.
-        """
-
-        checkpoint_path = os.path.join(self.checkpoint_dir, checkpoint_name)
-        try:
-            log(INFO, f"Saving state as {checkpoint_path}")
-            torch.save(checkpoint_dict, checkpoint_path)
-        except Exception as e:
-            log(ERROR, f"Encountered the following error while saving the checkpoint: {e}")
-            raise e
-
-    def load_checkpoint(self, checkpoint_name: str) -> dict[str, Any]:
-        """
-        Loads and returns the checkpoint stored in ``checkpoint_dir`` under the provided name if it exists. If it
-        does not exist, an assertion error will be thrown.
-
-        Args:
-            checkpoint_name (str): Name of the state checkpoint to be loaded.
-
-        Returns:
-            dict[str, Any]: A dictionary representing the checkpointed state, as loaded by ``torch.load``.
-        """
-
-        assert self.checkpoint_exists(checkpoint_name)
-        checkpoint_path = os.path.join(self.checkpoint_dir, checkpoint_name)
-        log(INFO, f"Loading state from checkpoint at {checkpoint_path}")
-
-        return torch.load(checkpoint_path)
-
-    def checkpoint_exists(self, checkpoint_name: str, **kwargs: Any) -> bool:
-        """
-        Checks if a checkpoint exists at the ``checkpoint_dir`` constructed at initialization + ``checkpoint_name``.
-
-        Args:
-            checkpoint_name (str): Name of checkpoint for existence test. Directory of checkpoint is held internally
-                as state by the checkpointer
-
-        Raises:
-            ValueError: Previously this function supported sending a path, but now requires ``checkpoint_name``. Will
-                raise an error is ``checkpoint_path`` provided.
-
-        Returns:
-            bool: True if checkpoint exists, otherwise false.
-        """
-        if "checkpoint_path" in kwargs:
-            raise ValueError(
-                "Previously this checkpoint supported sending a path, but it now requires a checkpoint_name only"
-            )
-        checkpoint_path = os.path.join(self.checkpoint_dir, checkpoint_name)
-        return os.path.exists(checkpoint_path)

fl4health/checkpointing/client_module.py

@@ -1,13 +1,19 @@
+from __future__ import annotations
+
 from collections.abc import Sequence
 from enum import Enum
 from logging import INFO
-from typing import Any
+from typing import TYPE_CHECKING
 
 import torch.nn as nn
 from flwr.common.logger import log
 from flwr.common.typing import Scalar
 
-from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer, TorchModuleCheckpointer
+if TYPE_CHECKING:
+    from fl4health.clients.basic_client import BasicClient
+
+from fl4health.checkpointing.checkpointer import TorchModuleCheckpointer
+from fl4health.checkpointing.state_checkpointer import ClientStateCheckpointer
 
 ModelCheckpointers = TorchModuleCheckpointer | Sequence[TorchModuleCheckpointer] | None
 
@@ -18,11 +24,12 @@ class CheckpointMode(Enum):
 
 
 class ClientCheckpointAndStateModule:
+
     def __init__(
         self,
         pre_aggregation: ModelCheckpointers = None,
         post_aggregation: ModelCheckpointers = None,
-        state_checkpointer: PerRoundStateCheckpointer | None = None,
+        state_checkpointer: ClientStateCheckpointer | None = None,
     ) -> None:
         """
         This module is meant to hold up three to major components that determine how clients handle model and state
@@ -44,9 +51,9 @@ def __init__(
             post_aggregation (ModelCheckpointers, optional): If defined, this checkpointer (or sequence
                 of checkpointers) is used to checkpoint models based on their validation metrics/losses **AFTER**
                 server-side aggregation. Defaults to None.
-            state_checkpointer (PerRoundStateCheckpointer | None, optional): If defined, this checkpointer is used to
-                preserve client state (not just models), in the event one wants to restart federated training.
-                Defaults to None.
+            state_checkpointer (ClientStateCheckpointer | None, optional): If defined, this checkpointer
+                is used to preserve client state (not just models), in the event one wants to restart
+                federated training. Defaults to None.
         """
         self.pre_aggregation = (
             [pre_aggregation] if isinstance(pre_aggregation, TorchModuleCheckpointer) else pre_aggregation
@@ -119,52 +126,42 @@ def maybe_checkpoint(
         else:
             raise ValueError(f"Unrecognized mode for checkpointing: {str(mode)}")
 
-    def save_state(self, state_checkpoint_name: str, state: dict[str, Any]) -> None:
+    def save_state(self, client: BasicClient) -> None:
         """
         This function is meant to facilitate saving state required to restart an FL process on the client side. This
-        function will simply save whatever information is passed in the state variable using the file name in
-        ``state_checkpoint_name``. This function should only be called if a ``state_checkpointer`` exists in this
-        module
+        function will simply save all the attributes stated in ``ClientStateCheckpointer.snapshot_attrs``.
+        This function should only be called if a ``state_checkpointer`` exists in this module.
 
         Args:
-            state_checkpoint_name (str): Name of the state checkpoint file. The checkpointer itself will have a
-                directory to which state will be saved.
-            state (dict[str, Any]): State to be saved so that training might be resumed on the client if federated
-                training is interrupted. For example, this might contain things like optimizer states, learning rate
-                scheduler states, etc.
+            client (BasicClient): The client object from which state will be saved.
 
         Raises:
             ValueError: Throws an error if this function is called, but no state checkpointer has been provided
         """
 
         if self.state_checkpointer is not None:
-            self.state_checkpointer.save_checkpoint(state_checkpoint_name, state)
+            self.state_checkpointer.save_client_state(client)
         else:
             raise ValueError("Attempting to save state but no state checkpointer is specified")
 
-    def maybe_load_state(self, state_checkpoint_name: str) -> dict[str, Any] | None:
+    def maybe_load_state(self, client: BasicClient) -> bool:
         """
-        This function facilitates loading of any pre-existing state (with the name ``state_checkpoint_name``) in the
-        directory of the ``state_checkpointer``. If the state already exists at the proper path, the state is loaded
-        and returned. If it doesn't exist, we return None.
+        This function facilitates loading of any pre-existing state (with the name ``checkpoint_name``) in the
+        directory of the ``checkpoint_dir``. If the state already exists at the proper path, the state is loaded
+        and will be automatically saved into client's attributes. If it doesn't exist, we return False.
 
         Args:
-            state_checkpoint_name (str): Name of the state checkpoint file. The checkpointer itself will have a
-                directory from which state will be loaded (if it exists).
+            client (BasicClient): client object into which state will be loaded if a checkpoint exists
 
         Raises:
             ValueError: Throws an error if this function is called, but no state checkpointer has been provided
 
         Returns:
-            dict[str, Any] | None: If the state checkpoint properly exists and is loaded correctly, this dictionary
-            carries that state. Otherwise, we return a None (or throw an exception).
+            bool : If the state checkpoint properly exists and is loaded correctly, client's attributes
+            are set to the loaded values, and True is returned. Otherwise, we return False (or throw an exception).
         """
 
         if self.state_checkpointer is not None:
-            if self.state_checkpointer.checkpoint_exists(state_checkpoint_name):
-                return self.state_checkpointer.load_checkpoint(state_checkpoint_name)
-            else:
-                log(INFO, "State checkpointer is defined but no state checkpoint exists.")
-                return None
+            return self.state_checkpointer.maybe_load_client_state(client)
         else:
             raise ValueError("Attempting to load state, but no state checkpointer is specified")