Merge pull request #232 from labmlai/archive

Archive

Author: lakshith-403

app/server/labml_app/db/__init__.py

@@ -11,7 +11,7 @@
 from pymongo.errors import ConnectionFailure
 
 from .. import settings
-from . import project
+from . import project, folder
 from . import user
 from . import status
 from . import app_token
@@ -58,6 +58,7 @@ def msave_dict(self, keys: List[str], data: List[ModelDict]):
 
 models = [user.User,
           project.Project,
+          folder.Folder,
           status.Status,
           status.RunStatus,
           app_token.AppToken,

app/server/labml_app/db/folder.py

@@ -0,0 +1,46 @@
+import time
+from enum import Enum
+from typing import List, Set
+
+from labml_db import Model
+
+from labml_app.db import run
+
+
+class DefaultFolders(Enum):
+    ARCHIVE = 'archive'
+    DEFAULT = 'default'
+
+
+class Folder(Model['Folder']):
+    name: str
+    created_at: float
+    run_uuids: Set[str]
+
+    @classmethod
+    def defaults(cls):
+        return dict(name='',
+                    created_at=time.time(),
+                    run_uuids=set(),
+                    )
+
+    def get_runs(self) -> List['run.Run']:
+        res = []
+        for run_uuid in self.run_uuids:
+            r_key = run.RunIndex.get(run_uuid)
+            if r_key is not None:
+                r = r_key.load()
+                if r is not None:
+                    res.append(r)
+
+        res.sort(key=lambda x: x.start_time, reverse=True)
+        return res
+
+    def add_run(self, run_uuid: str):
+        self.run_uuids.add(run_uuid)
+        self.save()
+
+    def remove_run(self, run_uuid: str):
+        self.run_uuids.remove(run_uuid)
+        self.save()
+

app/server/labml_app/db/project.py

@@ -2,7 +2,7 @@
 
 from labml_db import Model, Key, Index
 
-from . import run
+from . import run, folder
 from . import session
 from ..logger import logger
 
@@ -14,6 +14,7 @@ class Project(Model['Project']):
     runs: Dict[str, Key['run.Run']]
     sessions: Dict[str, Key['session.Session']]
     is_run_added: bool
+    folders: Dict[str, Key['folder.Folder']]
 
     @classmethod
     def defaults(cls):
@@ -23,6 +24,7 @@ def defaults(cls):
                     runs={},
                     sessions={},
                     is_run_added=False,
+                    folders={},
                     )
 
     def is_project_run(self, run_uuid: str) -> bool:
@@ -31,10 +33,15 @@ def is_project_run(self, run_uuid: str) -> bool:
     def is_project_session(self, session_uuid: str) -> bool:
         return session_uuid in self.sessions
 
-    def get_runs(self) -> List['run.Run']:
+    def get_runs(self, folder_name: str = "Default") -> List['run.Run']:
         res = []
         likely_deleted = []
-        for run_uuid, run_key in self.runs.items():
+        run_uuids = self.runs.keys()
+        if folder_name in self.folders:
+            f = self.folders[folder_name].load()
+            if f:
+                run_uuids = f.run_uuids
+        for run_uuid in run_uuids:
             try:
                 r = run.get(run_uuid)
                 if r:
@@ -71,6 +78,7 @@ def delete_runs(self, run_uuids: List[str], project_owner: str) -> None:
                 r = run.get(run_uuid)
                 if r and r.owner == project_owner:
                     try:
+                        self.delete_from_folder(r)
                         run.delete(run_uuid)
                     except TypeError:
                         logger.error(f'error while deleting the run {run_uuid}')
@@ -95,9 +103,24 @@ def add_run(self, run_uuid: str) -> None:
 
         if r:
             self.runs[run_uuid] = r.key
+            self.add_to_folder(folder.DefaultFolders.DEFAULT.value, r)
 
         self.save()
 
+    def add_run_with_model(self, r: run.Run) -> None:
+        self.runs[r.run_uuid] = r.key
+        self.is_run_added = True
+
+        self.add_to_folder(folder.DefaultFolders.DEFAULT.value, r)
+
+        self.save()
+
+    def get_run(self, run_uuid: str) -> Optional['run.Run']:
+        if run_uuid in self.runs:
+            return self.runs[run_uuid].load()
+        else:
+            return None
+
     def add_session(self, session_uuid: str) -> None:
         s = session.get(session_uuid)
 
@@ -106,6 +129,62 @@ def add_session(self, session_uuid: str) -> None:
 
         self.save()
 
+    def add_to_folder(self, folder_name: str, r: run.Run) -> None:
+        if folder_name not in self.folders:
+            f = folder.Folder(name=folder_name)
+            self.folders[folder_name] = f.key
+        else:
+            f = self.folders[folder_name].load()
+
+        f.add_run(r.run_uuid)
+        f.save()
+
+        r.parent_folder = f.name
+        r.save()
+
+    def remove_from_folder(self, folder_name: str, r: run.Run) -> None:
+        if folder_name not in self.folders:
+            f = folder.Folder(name=folder_name)
+            self.folders[folder_name] = f.key
+        else:
+            f = self.folders[folder_name].load()
+
+        f.remove_run(r.run_uuid)
+        f.save()
+
+        r.parent_folder = ''
+        r.save()
+
+    def archive_runs(self, run_uuids: List[str]) -> None:
+        for run_uuid in run_uuids:
+            if run_uuid in self.runs:
+                r = self.runs[run_uuid].load()
+                if r:
+                    self.remove_from_folder(folder.DefaultFolders.DEFAULT.value, r)
+                    self.add_to_folder(folder.DefaultFolders.ARCHIVE.value, r)
+
+        self.save()
+
+    def un_archive_runs(self, run_uuids: List[str]) -> None:
+        for run_uuid in run_uuids:
+            if run_uuid in self.runs:
+                r = self.runs[run_uuid].load()
+                if r:
+                    self.remove_from_folder(folder.DefaultFolders.ARCHIVE.value, r)
+                    self.add_to_folder(folder.DefaultFolders.DEFAULT.value, r)
+
+        self.save()
+
+    def delete_from_folder(self, r: run.Run) -> None:
+        folder_name = r.parent_folder
+        if folder_name in self.folders:
+            return
+        parent_folder = self.folders[folder_name].load()
+        if parent_folder is None:
+            return
+        parent_folder.remove_run(r.run_uuid)
+        parent_folder.save()
+
 
 class ProjectIndex(Index['Project']):
     pass
@@ -123,8 +202,8 @@ def get_project(labml_token: str) -> Union[None, Project]:
 def get_run(run_uuid: str, labml_token: str = '') -> Optional['run.Run']:
     p = get_project(labml_token)
 
-    if run_uuid in p.runs:
-        return p.runs[run_uuid].load()
+    if p.is_project_run(run_uuid):
+        return p.get_run(run_uuid)
     else:
         return None
 

app/server/labml_app/db/run.py

@@ -6,7 +6,7 @@
 from labml_db import Model, Key, Index, load_keys
 
 from .. import auth
-from . import user
+from . import user, folder
 from .. import utils
 from . import project
 from . import computer
@@ -64,6 +64,7 @@ class Run(Model['Run']):
     selected_configs: List['str']
     favourite_configs: List['str']
     main_rank: int
+    parent_folder: str
 
     wildcard_indicators: Dict[str, Dict[str, Union[str, bool]]]
     indicators: Dict[str, Dict[str, Union[str, bool]]]
@@ -110,6 +111,7 @@ def defaults(cls):
                     process_key=None,
                     session_id='',
                     main_rank=0,
+                    parent_folder='',
                     )
 
     @property
@@ -303,7 +305,8 @@ def get_data(self, request: Request, is_dist_run: bool = False) -> Dict[str, Uni
             'favourite_configs': self.favourite_configs,
             'selected_configs': self.selected_configs,
             'process_id': self.process_id,
-            'session_id': self.session_id
+            'session_id': self.session_id,
+            'folder': self.parent_folder
         }
 
     def get_summary(self) -> Dict[str, str]:
@@ -354,8 +357,8 @@ class RunIndex(Index['Run']):
 def get_or_create(request: Request, run_uuid: str, rank: int, world_size: int, main_rank: int, labml_token: str = '') -> 'Run':
     p = project.get_project(labml_token)
 
-    if run_uuid in p.runs:
-        return p.runs[run_uuid].load()
+    if p.is_project_run(run_uuid):
+        return p.get_run(run_uuid)
 
     run = get(run_uuid)
     if run is not None:
@@ -383,8 +386,7 @@ def get_or_create(request: Request, run_uuid: str, rank: int, world_size: int, m
               )
 
     if run.rank == 0:  # TODO
-        p.runs[run.run_uuid] = run.key
-        p.is_run_added = True
+        p.add_run_with_model(run)
 
     run.save()
     p.save()
@@ -414,10 +416,9 @@ def delete(run_uuid: str) -> None:
         analyses.AnalysisManager.delete_run(run_uuid)
 
 
-def get_runs(labml_token: str) -> List['Run']:
-    res = []
+def get_runs(labml_token: str, folder_name: str) -> List['Run']:
     p = project.get_project(labml_token)
-    load_keys(list(p.runs.values()))
+    res = p.get_runs(folder_name)
 
     return res
 

app/server/labml_app/handlers.py

@@ -10,7 +10,7 @@
 from .logger import logger
 from . import settings
 from . import auth
-from .db import run
+from .db import run, folder
 from .db import computer
 from .db import session
 from .db import user
@@ -214,12 +214,11 @@ async def claim_run(request: Request, run_uuid: str, token: Optional[str] = None
 
     default_project = u.default_project
 
-    if r.run_uuid not in default_project.runs:
+    if not default_project.is_project_run(run_uuid):
         # float_project = project.get_project(labml_token=settings.FLOAT_PROJECT_TOKEN)
 
         # if r.run_uuid in float_project.runs:
-        default_project.runs[r.run_uuid] = r.key
-        default_project.is_run_added = True
+        default_project.add_run_with_model(r)
         default_project.save()
         r.is_claimed = True
         r.owner = u.email
@@ -362,15 +361,16 @@ async def get_session_status(request: Request, session_uuid: str) -> JSONRespons
 
 @auth.login_required
 @auth.check_labml_token_permission
-async def get_runs(request: Request, labml_token: str, token: Optional[str] = None) -> EndPointRes:
+async def get_runs(request: Request, labml_token: str, token: Optional[str] = None,
+                   folder_name: str = folder.DefaultFolders.DEFAULT.value) -> EndPointRes:
     u = user.get_by_session_token(token)
 
     if labml_token:
-        runs_list = run.get_runs(labml_token)
+        runs_list = run.get_runs(labml_token, folder_name)
     else:
         default_project = u.default_project
         labml_token = default_project.labml_token
-        runs_list = default_project.get_runs()
+        runs_list = default_project.get_runs(folder_name)
 
     # run_uuids = [r.run_uuid for r in runs_list if r.world_size == 0]
     #
@@ -494,6 +494,39 @@ async def add_run(request: Request, run_uuid: str, token: Optional[str] = None)
     return {'is_successful': True}
 
 
+@auth.login_required
+async def archive_runs(request: Request, token: Optional[str] = None) -> EndPointRes:
+    json = await request.json()
+    run_uuids = json['run_uuids']
+
+    u = user.get_by_session_token(token)
+
+    try:
+        u.default_project.archive_runs(run_uuids)
+    except KeyError:
+        return {'is_successful': False, 'error': "Failed to archive. Probably due to inconsistencies with the server."
+                                                 "Please refresh the page and try again."}
+
+    return {'is_successful': True}
+
+
+@auth.login_required
+async def un_archive_runs(request: Request, token: Optional[str] = None) -> EndPointRes:
+    json = await request.json()
+    run_uuids = json['run_uuids']
+
+    u = user.get_by_session_token(token)
+
+    try:
+        u.default_project.un_archive_runs(run_uuids)
+    except KeyError:
+        return {'is_successful': False, 'error': "Failed to un-archive. Probably due to inconsistencies with the "
+                                                 "server."
+                                                 "Please refresh the page and try again."}
+
+    return {'is_successful': True}
+
+
 @auth.login_required
 async def add_session(request: Request, session_uuid: str, token: Optional[str] = None) -> EndPointRes:
     u = user.get_by_session_token(token)
@@ -553,6 +586,9 @@ def add_handlers(app: FastAPI):
     _add_server(app, 'POST', update_run, '{labml_token}/track')
     _add_server(app, 'POST', update_session, '{labml_token}/computer')
 
+    _add_ui(app, 'POST', archive_runs, 'runs/archive')
+    _add_ui(app, 'POST', un_archive_runs, 'runs/unarchive')
+
     _add_ui(app, 'GET', get_runs, 'runs/{labml_token}')
     _add_ui(app, 'PUT', delete_runs, 'runs')
     _add_ui(app, 'GET', get_sessions, 'sessions/{labml_token}')

app/ui/src/analyses/experiments/configs/view.ts

@@ -128,7 +128,7 @@ class RunConfigsView extends ScreenView {
             }
 
             await CACHE.getRun(this.uuid).updateRunData(data)
-            await CACHE.getRunsList().localUpdateRun(this.run)
+            await CACHE.getRunsList(this.run.folder).localUpdateRun(this.run)
         } catch (e) {
             this.userMessage.networkError(e, "Failed to save configurations")
             this.save.disabled = false