Repository: NVIDIA/NVFlare

Issue:
Issues on storing studying in Study Manager 



Discussion:

Resolving patch:
diff --git a/nvflare/apis/impl/study_manager.py b/nvflare/apis/impl/study_manager.py
--- a/nvflare/apis/impl/study_manager.py
+++ b/nvflare/apis/impl/study_manager.py
@@ -12,18 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import datetime
 import json
-import tempfile
+import logging
+from datetime import date, datetime
 from typing import List, Optional, Tuple
 
+from nvflare.apis.fl_context import FLContext
 from nvflare.apis.storage import StorageSpec
 from nvflare.apis.study_manager_spec import Study, StudyManagerSpec
 
 
 def custom_json_encoder(obj):
     try:
-        if isinstance(obj, datetime):
+        if isinstance(obj, (datetime, date)):
             return obj.isoformat()
         iterable = iter(obj)
     except TypeError:
@@ -33,29 +34,39 @@ def custom_json_encoder(obj):
 
 
 class StudyManager(StudyManagerSpec):
+    STORAGE_KEY = "std_mgr"
+
     def __init__(self, storage: StorageSpec):
+        self._name = self.__class__.__name__
+        self.logger = logging.getLogger(self._name)
         self._storage = storage
-        self._tmp_dir = tempfile.TemporaryDirectory()
-        self._existing_studies = dict()
+        try:
+            self._existing_studies = json.loads(self._storage.get_data(StudyManager.STORAGE_KEY).decode("utf-8"))
+        except BaseException:
+            self._existing_studies = list()
 
-    def create_study(self, study: Study) -> Tuple[Optional[Study], str]:
-        """Create the study object permanently
+    def add_study(self, study: Study, fl_ctx: FLContext) -> Tuple[Optional[Study], str]:
+        """Add the study object permanently
 
-        The caller must have validated the sites and users of the study.
+        The caller must have validated the participating_clients and participating_admins of the study.
 
         Validate the study before saving:
         The name of the study must be unique;
-        Sites and users must be defined;
-        Start and end time must make sense.
+        participating_clients and participating_admins must be defined;
+        Start and end date must make sense.
 
         Args:
             study: the caller-provided study info
+            fl_ctx: FLContext
+
+        Returns:
+            updated study info (e.g. created_at is set) and an emtpy string if successful
+            None and an error message if the provided study is not valid
 
-        Returns: updated study info (e.g. create_time is set)
 
         """
-        if self.get_study(study.name):
-            return None, "Unable to create duplicated study name."
+        if study.name in self._existing_studies:
+            return None, "Unable to add duplicated study name."
         if not study.participating_clients:
             return None, "Study has no participating clients."
         if not study.contact:
@@ -66,54 +77,100 @@ def create_study(self, study: Study) -> Tuple[Optional[Study], str]:
             return None, "Study has no start date."
         if not study.end_date:
             return None, "Study has no end date."
-        study.create_time = datetime.datetime.utcnow()
-        serialized_study = json.dumps(vars(study), default=custom_json_encoder).encode("utf-8")
-        self._store.create_object(
-            study.name, serialized_study, {"start_date": study.start_date, "end_date": study.end_date}
-        )
-        deserialized_study = json.loads(serialized_study.decode("utf-8"))
-        deserialized_study["start_date"] = datetime.date.fromisoformat(deserialized_study["start_date"])
-        deserialized_study["end_date"] = datetime.date.fromisoformat(deserialized_study["end_date"])
-        stored_study = Study(**deserialized_study)  # for debugging purpose
+        if study.end_date < study.start_date:
+            return (
+                None,
+                "Expect end_date later than start_date.  Got start_date={study.start_date} and end_data={study.end_date}",
+            )
+        study.created_at = datetime.utcnow()
+        serialized_study = self._study_to_bytes(study)
+        try:
+            self._storage.create_object(
+                study.name, serialized_study, {"start_date": study.start_date, "end_date": study.end_date}
+            )
+            self._storage.create_object(
+                StudyManager.STORAGE_KEY,
+                json.dumps(self._existing_studies + [study.name]).encode("utf-8"),
+                meta={},
+                overwrite_existing=True,
+            )
+        except BaseException as e:
+            self.logger.warning(f"Unable to add study into storage due to {e}")
+            return None, f"Unable to add study into storage due to {e}"
+        self._existing_studies.append(study.name)
         return study, ""
 
-    def list_studies(self) -> List[str]:
+    @staticmethod
+    def _study_from_bytes(bytes):
+        try:
+            deserialized_study = json.loads(bytes.decode("utf-8"))
+            deserialized_study["start_date"] = date.fromisoformat(deserialized_study.pop("start_date"))
+            deserialized_study["end_date"] = date.fromisoformat(deserialized_study.pop("end_date"))
+            deserialized_study["created_at"] = datetime.fromisoformat(deserialized_study.pop("created_at"))
+            return Study(**deserialized_study)
+        except BaseException as e:
+            return None
+
+    def _study_to_bytes(self, study):
+        try:
+            serialized_study = json.dumps(vars(study), default=custom_json_encoder).encode("utf-8")
+            return serialized_study
+        except BaseException:
+            return None
+
+    def list_studies(self, fl_ctx: FLContext) -> List[str]:
         """
         List names of all defined studies
 
+        Args:
+            fl_ctx: FLContext
+
         Returns: list of study names
 
         """
         # either self._storage.list_object() or
-        return list(self._existing_studies)
+        return self._existing_studies
 
-    def list_active_studies(self) -> List[str]:
+    def list_active_studies(self, fl_ctx: FLContext) -> List[str]:
         """
         List names of all active studies (started but not ended)
 
+        Args:
+            fl_ctx: FLContext
+
         Returns: list of study names
 
         """
-        now = datetime.datetime.utcnow()
+        today = date.today()
         active_studies = list()
-        for k, v in self._existing_studies:
-            if v["start_date"] >= now and v["end_date"] <= now:
-                active_studies.append(k)
+        for st in self._existing_studies:
+            try:
+                meta = self._storage.get_meta(st)
+                if meta["start_date"] > today:
+                    continue
+                if meta["end_date"] <= today:
+                    continue
+                active_studies.append(st)
+            except BaseException as e:
+                self.logger.warning(
+                    f"Inconsistence between internal persisted states.  Expect {st} with meta in storage but got {e}."
+                )
         return active_studies
 
-    def get_study(self, name: str) -> Study:
+    def get_study(self, name: str, fl_ctx: FLContext) -> Optional[Study]:
         """Get the Study object for the specified name.
 
         Args:
             name: unique name of the study
+            fl_ctx: FLContext
 
         Returns: the Study object
 
         """
-        serialized_study = self._storage.get_data(name)
-        if serialized_study:
-            deserialized_study = json.loads(serialized_study.decode("utf-8"))
-        deserialized_study["start_date"] = datetime.date.fromisoformat(deserialized_study["start_date"])
-        deserialized_study["end_date"] = datetime.date.fromisoformat(deserialized_study["end_date"])
-        study = Study(**deserialized_study)
+        try:
+            serialized_study = self._storage.get_data(name)
+        except RuntimeError as e:
+            self.logger.warning(f"Unable to load study {name} from storage.")
+            return None
+        study = self._study_from_bytes(serialized_study)
         return study
diff --git a/nvflare/apis/study_manager_spec.py b/nvflare/apis/study_manager_spec.py
--- a/nvflare/apis/study_manager_spec.py
+++ b/nvflare/apis/study_manager_spec.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from datetime import datetime
+from datetime import date, datetime
 from typing import List
 
 from .fl_context import FLContext
@@ -27,9 +27,10 @@ def __init__(
         contact: str,
         participating_clients: List[str],
         participating_admins: List[str],
-        start_date: datetime.date,
-        end_date: datetime.date,
+        start_date: date,
+        end_date: date,
         reviewers=None,
+        created_at: datetime = None,
     ):
         self.name = name
         self.description = description
@@ -39,7 +40,7 @@ def __init__(
         self.start_date = start_date
         self.end_date = end_date
         self.reviewers = reviewers
-        self.created_at = datetime.utcnow().isoformat()
+        self.created_at = created_at
 
 
 class StudyManagerSpec(ABC):
@@ -57,7 +58,10 @@ def add_study(self, study: Study, fl_ctx: FLContext) -> Study:
         Args:
             study: the caller-provided study info
 
-        Returns: updated study info (e.g. create_time is set)
+        Returns:
+            updated study info (e.g. created_at is set) and an emtpy string if successful
+            None and an error message if the provided study is not valid
+
 
         """
         pass


Test patch: