Apply EAV data model to system attributes.

pfnopt/storages/base.py

@@ -11,8 +11,6 @@
 from pfnopt import distributions  # NOQA
 from pfnopt import structs  # NOQA
 
-SYSTEM_ATTRS_KEY = '__system__'
-
 
 @six.add_metaclass(abc.ABCMeta)
 class BaseStorage(object):
@@ -37,14 +35,11 @@ def set_study_task(self, study_id, task):
 
         raise NotImplementedError
 
-    # TODO(sano): support setting attribute in a thread-safe way.
+    @abc.abstractmethod
     def set_study_system_attr(self, study_id, key, value):
         # type: (int, str, Any) -> None
 
-        user_attrs = self.get_study_user_attrs(study_id)
-        user_attrs[SYSTEM_ATTRS_KEY][key] = value
-        self.set_study_user_attr(
-            study_id, SYSTEM_ATTRS_KEY, user_attrs[SYSTEM_ATTRS_KEY])
+        raise NotImplementedError
 
     # Basic study access
 
@@ -72,11 +67,11 @@ def get_study_user_attrs(self, study_id):
 
         raise NotImplementedError
 
+    @abc.abstractmethod
     def get_study_system_attr(self, study_id, key):
         # type: (int, str) -> Any
 
-        user_attrs = self.get_study_user_attrs(study_id)
-        return copy.deepcopy(user_attrs[SYSTEM_ATTRS_KEY][key])
+        raise NotImplementedError
 
     @abc.abstractmethod
     def get_all_study_summaries(self):
@@ -128,14 +123,11 @@ def set_trial_user_attr(self, trial_id, key, value):
 
         raise NotImplementedError
 
-    # TODO(sano): support setting attribute in a thread-safe way.
+    @abc.abstractmethod
     def set_trial_system_attr(self, trial_id, key, value):
         # type: (int, str, Any) -> None
 
-        user_attrs = self.get_trial(trial_id).user_attrs
-        user_attrs[SYSTEM_ATTRS_KEY][key] = value
-        self.set_trial_user_attr(
-            trial_id, SYSTEM_ATTRS_KEY, user_attrs[SYSTEM_ATTRS_KEY])
+        raise NotImplementedError
 
     # Basic trial access
 
@@ -182,8 +174,7 @@ def get_trial_user_attrs(self, trial_id):
     def get_trial_system_attr(self, trial_id, key):
         # type: (int, str) -> Any
 
-        user_attrs = self.get_trial(trial_id).user_attrs
-        return copy.deepcopy(user_attrs[SYSTEM_ATTRS_KEY][key])
+        return copy.deepcopy(self.get_trial(trial_id).system_attrs[key])
 
     # Methods for the TPE sampler
 

pfnopt/storages/in_memory.py

@@ -23,6 +23,7 @@ def __init__(self):
         self.param_distribution = {}  # type: Dict[str, distributions.BaseDistribution]
         self.task = structs.StudyTask.NOT_SET
         self.study_user_attrs = {}  # type: Dict[str, Any]
+        self.study_system_attrs = {}  # type: Dict[str, Any]
 
         self._lock = threading.Lock()
 
@@ -40,8 +41,6 @@ def __setstate__(self, state):
     def create_new_study_id(self):
         # type: () -> int
 
-        self.study_user_attrs[base.SYSTEM_ATTRS_KEY] = {}
-
         return IN_MEMORY_STORAGE_STUDY_ID  # TODO(akiba)
 
     def set_study_task(self, study_id, task):
@@ -59,6 +58,12 @@ def set_study_user_attr(self, study_id, key, value):
         with self._lock:
             self.study_user_attrs[key] = value
 
+    def set_study_system_attr(self, study_id, key, value):
+        # type: (int, str, Any) -> None
+
+        with self._lock:
+            self.study_system_attrs[key] = value
+
     def get_study_id_from_uuid(self, study_uuid):
         # type: (str) -> int
 
@@ -82,6 +87,16 @@ def get_study_user_attrs(self, study_id):
         with self._lock:
             return copy.deepcopy(self.study_user_attrs)
 
+    def get_study_system_attr(self, study_id, key):
+        # type: (int, str) -> Any
+
+        with self._lock:
+            try:
+                return copy.deepcopy(self.study_system_attrs[key])
+            except KeyError:
+                raise ValueError(
+                    'System attribute {} does not exist in Study {}.'.format(key, study_id))
+
     def get_all_study_summaries(self):
         # type: () -> List[structs.StudySummary]
 
@@ -100,6 +115,7 @@ def get_all_study_summaries(self):
             task=self.task,
             best_trial=best_trial,
             user_attrs=copy.deepcopy(self.study_user_attrs),
+            system_attrs=copy.deepcopy(self.study_system_attrs),
             n_trials=len(self.trials),
             datetime_start=datetime_start
         )]
@@ -115,7 +131,8 @@ def create_new_trial_id(self, study_id):
                     trial_id=trial_id,
                     state=structs.TrialState.RUNNING,
                     params={},
-                    user_attrs={base.SYSTEM_ATTRS_KEY: {}},
+                    user_attrs={},
+                    system_attrs={},
                     value=None,
                     intermediate_values={},
                     params_in_internal_repr={},
@@ -186,6 +203,12 @@ def set_trial_user_attr(self, trial_id, key, value):
         with self._lock:
             self.trials[trial_id].user_attrs[key] = value
 
+    def set_trial_system_attr(self, trial_id, key, value):
+        # type: (int, str, Any) -> None
+
+        with self._lock:
+            self.trials[trial_id].system_attrs[key] = value
+
     def get_trial(self, trial_id):
         # type: (int) -> structs.FrozenTrial
 

pfnopt/storages/rdb/models.py

@@ -18,7 +18,7 @@
 from pfnopt.structs import StudyTask
 from pfnopt.structs import TrialState
 
-SCHEMA_VERSION = 6
+SCHEMA_VERSION = 8
 MAX_ATTR_LENGTH = 2048
 NOT_FOUND_MSG = 'Record does not exist.'
 
@@ -79,7 +79,7 @@ class StudyUserAttributeModel(BaseModel):
     __table_args__ = (UniqueConstraint('study_id', 'key'), )  # type: Any
     study_user_attribute_id = Column(Integer, primary_key=True)
     study_id = Column(Integer, ForeignKey('studies.study_id'))
-    key = Column(String(255))
+    key = Column(String(MAX_ATTR_LENGTH))
     value_json = Column(String(MAX_ATTR_LENGTH))
 
     study = orm.relationship(StudyModel)
@@ -100,13 +100,38 @@ def where_study_id(cls, study_id, session):
         return session.query(cls).filter(cls.study_id == study_id).all()
 
 
+class StudySystemAttributeModel(BaseModel):
+    __tablename__ = 'study_system_attributes'
+    __table_args__ = (UniqueConstraint('study_id', 'key'), )  # type: Any
+    study_system_attribute_id = Column(Integer, primary_key=True)
+    study_id = Column(Integer, ForeignKey('studies.study_id'))
+    key = Column(String(MAX_ATTR_LENGTH))
+    value_json = Column(String(MAX_ATTR_LENGTH))
+
+    study = orm.relationship(StudyModel)
+
+    @classmethod
+    def find_by_study_and_key(cls, study, key, session):
+        # type: (StudyModel, str, orm.Session) -> Optional[StudySystemAttributeModel]
+
+        attribute = session.query(cls). \
+            filter(cls.study_id == study.study_id).filter(cls.key == key).one_or_none()
+
+        return attribute
+
+    @classmethod
+    def where_study_id(cls, study_id, session):
+        # type: (int, orm.Session) -> List[StudySystemAttributeModel]
+
+        return session.query(cls).filter(cls.study_id == study_id).all()
+
+
 class TrialModel(BaseModel):
     __tablename__ = 'trials'
     trial_id = Column(Integer, primary_key=True)
     study_id = Column(Integer, ForeignKey('studies.study_id'))
     state = Column(Enum(TrialState), nullable=False)
     value = Column(Float)
-    user_attributes_json = Column(String(MAX_ATTR_LENGTH))
     datetime_start = Column(DateTime, default=datetime.now)
     datetime_complete = Column(DateTime)
 
@@ -157,6 +182,88 @@ def all(cls, session):
         return session.query(cls).all()
 
 
+class TrialUserAttributeModel(BaseModel):
+    __tablename__ = 'trial_user_attributes'
+    __table_args__ = (UniqueConstraint('trial_id', 'key'), )  # type: Any
+    trial_user_attribute_id = Column(Integer, primary_key=True)
+    trial_id = Column(Integer, ForeignKey('trials.trial_id'))
+    key = Column(String(MAX_ATTR_LENGTH))
+    value_json = Column(String(MAX_ATTR_LENGTH))
+
+    trial = orm.relationship(TrialModel)
+
+    @classmethod
+    def find_by_trial_and_key(cls, trial, key, session):
+        # type: (TrialModel, str, orm.Session) -> Optional[TrialUserAttributeModel]
+
+        attribute = session.query(cls). \
+            filter(cls.trial_id == trial.trial_id).filter(cls.key == key).one_or_none()
+
+        return attribute
+
+    @classmethod
+    def where_study(cls, study, session):
+        # type: (StudyModel, orm.Session) -> List[TrialUserAttributeModel]
+
+        trial_user_attributes = session.query(cls).join(TrialModel). \
+            filter(TrialModel.study_id == study.study_id).all()
+
+        return trial_user_attributes
+
+    @classmethod
+    def where_trial(cls, trial, session):
+        # type: (TrialModel, orm.Session) -> List[TrialUserAttributeModel]
+
+        return session.query(cls).filter(cls.trial_id == trial.trial_id).all()
+
+    @classmethod
+    def all(cls, session):
+        # type: (orm.Session) -> List[TrialUserAttributeModel]
+
+        return session.query(cls).all()
+
+
+class TrialSystemAttributeModel(BaseModel):
+    __tablename__ = 'trial_system_attributes'
+    __table_args__ = (UniqueConstraint('trial_id', 'key'), )  # type: Any
+    trial_system_attribute_id = Column(Integer, primary_key=True)
+    trial_id = Column(Integer, ForeignKey('trials.trial_id'))
+    key = Column(String(MAX_ATTR_LENGTH))
+    value_json = Column(String(MAX_ATTR_LENGTH))
+
+    trial = orm.relationship(TrialModel)
+
+    @classmethod
+    def find_by_trial_and_key(cls, trial, key, session):
+        # type: (TrialModel, str, orm.Session) -> Optional[TrialSystemAttributeModel]
+
+        attribute = session.query(cls). \
+            filter(cls.trial_id == trial.trial_id).filter(cls.key == key).one_or_none()
+
+        return attribute
+
+    @classmethod
+    def where_study(cls, study, session):
+        # type: (StudyModel, orm.Session) -> List[TrialSystemAttributeModel]
+
+        trial_system_attributes = session.query(cls).join(TrialModel). \
+            filter(TrialModel.study_id == study.study_id).all()
+
+        return trial_system_attributes
+
+    @classmethod
+    def where_trial(cls, trial, session):
+        # type: (TrialModel, orm.Session) -> List[TrialSystemAttributeModel]
+
+        return session.query(cls).filter(cls.trial_id == trial.trial_id).all()
+
+    @classmethod
+    def all(cls, session):
+        # type: (orm.Session) -> List[TrialSystemAttributeModel]
+
+        return session.query(cls).all()
+
+
 class TrialParamModel(BaseModel):
     __tablename__ = 'trial_params'
     __table_args__ = (UniqueConstraint('trial_id', 'param_name'), )  # type: Any