[feat] Add ability to define properties on Container subclasses

pkgs/aimstack/asp/models/run.py

@@ -6,7 +6,7 @@
 
 from functools import partialmethod
 
-from aim import Container
+from aim import Container, Property
 from aim._sdk.utils import utc_timestamp
 from aim._sdk import type_utils
 from aimcore.callbacks import Caller
@@ -41,6 +41,10 @@
 
 @type_utils.query_alias('run')
 class Run(Container, Caller):
+    name = Property()
+    description = Property(default='')
+    archived = Property(default=False)
+
     def __init__(self, hash_: Optional[str] = None, *,
                  repo: Optional[Union[str, 'Repo']] = None,
                  mode: Optional[Union[str, ContainerOpenMode]] = ContainerOpenMode.WRITE):
@@ -49,10 +53,6 @@ def __init__(self, hash_: Optional[str] = None, *,
         if not self._is_readonly:
             if self.name is None:
                 self.name = f'Run #{self.hash}'
-            if self.description is None:
-                self.description = ''
-            if self.archived is None:
-                self.archived = False
 
     def enable_system_monitoring(self):
         if not self._is_readonly:
@@ -81,30 +81,6 @@ def track_system_resources(self, stats: Dict[str, Any], context: Dict, **kwargs)
         for resource_name, usage in stats.items():
             self.sequences.typed_sequence(SystemMetric, resource_name, context).track(usage)
 
-    @property
-    def name(self) -> str:
-        return self._attrs_tree.get('name', None)
-
-    @name.setter
-    def name(self, val: str):
-        self['name'] = val
-
-    @property
-    def description(self) -> str:
-        return self._attrs_tree.get('description', None)
-
-    @description.setter
-    def description(self, val: str):
-        self['description'] = val
-
-    @property
-    def archived(self) -> bool:
-        return self._attrs_tree.get('archived', None)
-
-    @archived.setter
-    def archived(self, val: bool):
-        self['archived'] = val
-
     @property
     def creation_time(self) -> float:
         return self._tree[KeyNames.INFO_PREFIX, 'creation_time']

src/aimcore/web/ui/public/aim_ui_core.py

@@ -52,6 +52,22 @@ class WaitForQueryError(Exception):
     pass
 
 
+# dict forbids setting attributes, hence using derived class
+class dictionary(dict):
+    pass
+
+
+def process_properties(obj: dict):
+    if '$properties' in obj:
+        props = obj.pop('$properties')
+        new_obj = dictionary()
+        new_obj.update(obj)
+        for k, v in props.items():
+            setattr(new_obj, k, v)
+        return new_obj
+    return obj
+
+
 def query_filter(type_, query="", count=None, start=None, stop=None, is_sequence=False):
     query_key = f"{type_}_{query}_{count}_{start}_{stop}"
 
@@ -60,7 +76,7 @@ def query_filter(type_, query="", count=None, start=None, stop=None, is_sequence
 
     try:
         data = search(board_path, type_, query, count, start, stop, is_sequence)
-        data = json.loads(data)
+        data = json.loads(data, object_hook=process_properties)
 
         query_results_cache[query_key] = data
         return data
@@ -79,7 +95,7 @@ def run_function(func_name, params):
 
     try:
         res = runFunction(board_path, func_name, params)
-        data = json.loads(res)["value"]
+        data = json.loads(res, object_hook=process_properties)["value"]
 
         query_results_cache[run_function_key] = data
         return data

src/python/aim/__init__.py

@@ -2,7 +2,7 @@
 
 from .record import Record
 from .sequence import Sequence
-from .container import Container
+from .container import Container, Property
 from .repo import Repo
 
 from aim._ext.notebook.notebook import load_ipython_extension

src/python/aim/_sdk/container.py

@@ -78,6 +78,28 @@ def _close(self) -> None:
             self._lock.release()
 
 
+class Property:
+    PROP_NAME_BLACKLIST = (  # do not allow property names to be dict class public methods
+        'clear', 'copy', 'fromkeys', 'get', 'items', 'keys', 'pop', 'popitem', 'setdefault', 'update', 'values')
+
+    def __init__(self, default=None):
+        self._default = default
+        self._name = None  # Will be set by __set_name__
+
+    def __set_name__(self, owner, name):
+        if name in Property.PROP_NAME_BLACKLIST:
+            raise RuntimeError(f'Cannot define Aim Property with name \'{name}\'.')
+        self._name = name
+
+    def __get__(self, instance: 'Container', owner):
+        if instance is None:
+            return self
+        return instance._get_property(self._name, self._default)
+
+    def __set__(self, instance: 'Container', value: Any):
+        instance._set_property(self._name, value)
+
+
 @type_utils.query_alias('container', 'c')
 @type_utils.auto_registry
 class Container(ABCContainer):
@@ -138,6 +160,10 @@ def __init__(self, hash_: Optional[str] = None, *,
                 self._meta_tree[KeyNames.CONTAINERS, self.get_typename()] = 1
                 self[...] = {}
 
+                for attr_name, attr in self.__class__.__dict__.items():
+                    if isinstance(attr, Property):
+                        self._set_property(attr_name, attr._default)
+
             self._tree[KeyNames.INFO_PREFIX, 'end_time'] = None
 
         self._resources = ContainerAutoClean(self)
@@ -171,6 +197,9 @@ def __storage_init__(self):
         self._meta_attrs_tree: TreeView = self._meta_tree.subtree('attrs')
         self._attrs_tree: TreeView = self._tree.subtree('attrs')
 
+        self._meta_props_tree: TreeView = self._meta_tree.subtree('_props')
+        self._props_tree: TreeView = self._tree.subtree('_props')
+
         self._data_loader: Callable[[], 'TreeView'] = lambda: self._sequence_data_tree
         self.__sequence_data_tree: TreeView = None
         self._sequence_map = ContainerSequenceMap(self, Sequence)
@@ -202,6 +231,16 @@ def get(self, key, default: Any = None, strict: bool = False):
         except KeyError:
             return default
 
+    def _set_property(self, name: str, value: Any):
+        self._props_tree[name] = value
+        self._meta_props_tree.merge(name, value)
+
+    def _get_property(self, name: str, default: Any = None) -> Any:
+        return self._props_tree.get(name, default)
+
+    def collect_properties(self) -> Dict:
+        return self._props_tree.collect()
+
     def match(self, expr) -> bool:
         query = RestrictedPythonQuery(expr)
         query_cache = {}

src/python/aim/container.py

@@ -1 +1 @@
-from aim._sdk.container import Container  # noqa F401
+from aim._sdk.container import Container, Property  # noqa F401

src/python/aim/utils.py

@@ -95,12 +95,14 @@ def sequence_data(
 
 
 def container_data(container: Container) -> Dict:
-    data = {
+    data = container[...]
+    data.update({
         'hash': container.hash,
         'params': container[...],
+        '$properties': container.collect_properties(),
         'container_type': container.get_typename(),
         'container_full_type': container.get_full_typename(),
-    }
+    })
     return data