NodeStateManager

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

    Args:
        job_id: job_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,
        "job_id": job_id
    }

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

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

    Args:
        job_id: job_id
        round_nb: current Round 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, "job_id_%s" % job_id)
    try:
        os.makedirs(base_dir, exist_ok=True)

    except Exception as e:
        raise FedbiomedNodeStateManagerError(f"{ErrorNumbers.FB323.value}: Failing to create directories "
                                             f"{base_dir}") from e
    # TODO catch exception here
    file_name = element.value % (round_nb, self._state_id)
    return os.path.join(base_dir, file_name)

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

    Args:
        job_id: the job for which a state is requested
        state_id: the unique identifier of the job

    Returns:
        dict containing the job state

    Raises:
        FedbiomedNodeStateManagerError: no matching state in the database
    """
    state = self._load_state(job_id, state_id)

    if state is None:
        raise FedbiomedNodeStateManagerError(f"{ErrorNumbers.FB323.value}: no entries matching job_id {job_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 whch to reload a Node state
        testing: only doing testing, not training
    """

    self._previous_state_id = previous_state_id
    if not testing:
        self._generate_new_state_id()

        self._node_state_base_dir = os.path.join(environ["VAR_DIR"], "node_state_%s" % environ["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, job_id: str):
    raise NotImplementedError

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