NodeStateManager

def add(self, experiment_id: str, state: Dict[str, Dict[str, str]]) -> str:
    """Adds new `Node` State into Database.

    Args:
        experiment_id: experiment_id used to save entry in database
        state: state to be saved in database.

    Returns:
        state_id

    Raises:
        FedbiomedNodeStateManagerError: state manager is not initialized
    """

    if self._state_id is None:
        raise FedbiomedNodeStateManagerError(
            f"{ErrorNumbers.FB323.value}: Node state manager is not initialized"
        )
    header = {
        "version_node_id": str(__node_state_version__),
        "state_id": self._state_id,
        "experiment_id": experiment_id,
    }

    state.update(header)
    self._save_state(self._state_id, state)
    return self._state_id

def generate_folder_and_create_file_name(
    self, experiment_id: str, sequence_nb: int, element: NodeStateFileName
) -> str:
    """Creates folders and file name for each content (Optimizer, model, ...) that composes
    a Node State.

    Args:
        experiment_id: experiment_id
        sequence_nb: current Round number ou Step Number
        element: a NodeStateFileName object used to create specific file names. For instance,
            could be a NodeStateFileName.OPTIMIZER

    Raises:
        FedbiomedNodeStateManagerError: raised if node state manager is not initialized
        FedbiomedNodeStateManagerError: raised if folder can not be created.

    Returns:
        path to the file that corresponds to the object that needs to be saved.
    """

    node_state_base_dir = self.get_node_state_base_dir()
    if node_state_base_dir is None:
        raise FedbiomedNodeStateManagerError(
            f"{ErrorNumbers.FB323.value}: working directory has not been "
            "initialized, have you run `initialize` method beforehand ?"
        )

    if self._state_id is None:
        raise FedbiomedNodeStateManagerError(
            f"{ErrorNumbers.FB323.value}: Node state manager is not initialized"
        )
    base_dir = os.path.join(node_state_base_dir, "experiment_id_%s" % experiment_id)
    try:
        os.makedirs(base_dir, exist_ok=True)

    except Exception as e:
        raise FedbiomedNodeStateManagerError(
            f"{ErrorNumbers.FB323.value}: Failing to create directories {base_dir}"
        ) from e
    # TODO catch exception here
    file_name = element.value % (sequence_nb, self._state_id)
    path = os.path.join(base_dir, file_name)
    logger.debug(
        f"Generated node state file path: experiment_id={experiment_id} "
        f"round={sequence_nb} element={getattr(element, 'name', repr(element))} path={path}"
    )
    return path

def get(self, experiment_id: str, state_id: str) -> Dict:
    """Returns a state of an experiment on the `Node`.

    Args:
        experiment_id: the experiment for which a state is requested
        state_id: the unique identifier of the experiment

    Returns:
        dict containing the experiment state

    Raises:
        FedbiomedNodeStateManagerError: no matching state in the database
    """
    logger.debug(
        f"Loading node state: experiment_id={experiment_id} state_id={state_id}"
    )
    state = self._load_state(experiment_id, state_id)

    if state is None:
        raise FedbiomedNodeStateManagerError(
            f"{ErrorNumbers.FB323.value}: no entries matching"
            f" experiment_id {experiment_id} and "
            f"state_id {state_id} found in the DataBase"
        )

    # from this point, state should be a dictionary
    self._check_version(state.get("version_node_id"))
    return state

def get_node_state_base_dir(self) -> Optional[str]:
    """Returns `Node` State base directory path, in which are saved Node state files and other contents

    Returns:
        path to `Node` state base directory, or None if not defined
    """
    return self._node_state_base_dir

def initialize(
    self, previous_state_id: Optional[str] = None, testing: Optional[bool] = False
) -> None:
    """Initializes NodeStateManager, by creating folder that will contains Node state folders.

    Args:
        previous_state_id: state_id from previous Round, from which to reload a Node state
        testing: only doing testing, not training

    Raises:
        FedbiomedNodeStateManagerError: raised if needed directories cannot be created
    """

    self._previous_state_id = previous_state_id
    logger.debug(
        f"Initializing node state manager: node_id={self._node_id} "
        f"previous_state_id={previous_state_id} testing={testing}"
    )
    if not testing:
        self._generate_new_state_id()

        self._node_state_base_dir = os.path.join(
            self._dir, VAR_FOLDER_NAME, "node_state_%s" % self._node_id
        )
        # Always create the base folder for saving states for this node
        try:
            os.makedirs(self._node_state_base_dir, exist_ok=True)
        except Exception as e:
            raise FedbiomedNodeStateManagerError(
                f"{ErrorNumbers.FB323.value}: Failing to create"
                f" directories {self._node_state_base_dir}"
            ) from e

def list_states(self, experiment_id: str):
    raise NotImplementedError

def remove(self, experiment_id: Optional[str], state_id: Optional[str]):
    raise NotImplementedError