Merge pull request #307 from labmlai/data-store

Data store

app/local_server_settings/server/analyses_settings.py

@@ -15,6 +15,8 @@
 from .analyses.experiments.stderr import StdErrModel
 from .analyses.experiments.stdlogger import StdLoggerModel
 
+from .analyses.experiments.data_store import DataStoreModel
+
 experiment_analyses = [MetricsAnalysis]
 
 computer_analyses = [CPUAnalysis,

app/server/labml_app/analyses/experiments/data_store.py

@@ -0,0 +1,59 @@
+from typing import Dict, Any
+
+from labml_db import Index, Model
+from labml_db.serializer.pickle import PickleSerializer
+from labml_db.serializer.yaml import YamlSerializer
+
+from labml_app.analyses.analysis import Analysis
+
+
+@Analysis.db_model(PickleSerializer, 'data_store')
+class DataStoreModel(Model['data_store']):
+    data: Dict[str, Any]
+
+    @classmethod
+    def defaults(cls):
+        return dict(
+            data={}
+        )
+
+    def get_data(self):
+        return self.data
+
+    def update_data(self, data: Dict[str, Any]):
+        for key, value in data.items():
+            self.data[key] = value
+
+
+@Analysis.db_index(YamlSerializer, 'datastore_index.yaml')
+class DataStoreIndex(Index['data_store']):
+    pass
+
+
+@Analysis.route('GET', 'datastore/{run_uuid}')
+async def get_data_store(run_uuid: str) -> Any:
+    key = DataStoreIndex.get(run_uuid)
+    if key is None:
+        data_store = DataStoreModel()
+        data_store.save()
+        DataStoreIndex.set(run_uuid, data_store.key)
+
+        return data_store.get_data()
+    else:
+        return key.load().get_data()
+
+
+@Analysis.route('POST', 'datastore/{run_uuid}')
+async def update_data_store(run_uuid: str, data: Dict[str, Any]) -> Any:
+    key = DataStoreIndex.get(run_uuid)
+    if key is None:
+        data_store = DataStoreModel()
+        data_store.save()
+        DataStoreIndex.set(run_uuid, data_store.key)
+    else:
+        data_store = key.load()
+
+    data_store.update_data(data)
+    data_store.save()
+
+    return data_store.get_data()

app/server/labml_app/analyses_settings.sample.py

@@ -14,6 +14,8 @@
 from .analyses.experiments.stderr import StdErrModel
 from .analyses.experiments.stdlogger import StdLoggerModel
 
+from .analyses.experiments.data_store import DataStoreModel
+
 experiment_analyses = [MetricsAnalysis]
 
 computer_analyses = [CPUAnalysis,

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

@@ -1,26 +1,20 @@
 import {Weya as $, WeyaElement} from "../../../../../lib/weya/weya"
 import {ROUTER, SCREEN} from "../../../app"
 import {Run} from "../../../models/run"
-import CACHE, {RunCache, RunsListCache, RunStatusCache, UserCache} from "../../../cache/cache"
+import CACHE, {DataStoreCache, RunCache, RunStatusCache, UserCache} from "../../../cache/cache"
 import {Status} from "../../../models/status"
-import {
-    BackButton,
-    CancelButton,
-    DeleteButton,
-    EditButton,
-    SaveButton,
-} from "../../../components/buttons"
+import {BackButton, DeleteButton, SaveButton,} from "../../../components/buttons"
 import EditableField from "../../../components/input/editable_field"
 import {formatTime, getTimeDiff} from "../../../utils/time"
 import {DataLoader} from "../../../components/loader"
-import {BadgeView} from "../../../components/badge"
 import {StatusView} from "../../../components/status"
-import {handleNetworkError, handleNetworkErrorInplace} from '../../../utils/redirect'
+import {handleNetworkErrorInplace} from '../../../utils/redirect'
 import {setTitle} from '../../../utils/document'
 import {formatFixed} from "../../../utils/value"
 import {ScreenView} from '../../../screen_view'
 import {User} from '../../../models/user'
 import {UserMessages} from "../../../components/user_messages"
+import {DataStore} from "../../../models/data_store";
 
 enum EditStatus {
     NOCHANGE,
@@ -36,6 +30,8 @@ class RunHeaderView extends ScreenView {
     statusCache: RunStatusCache
     user: User
     userCache: UserCache
+    dataStoreCache: DataStoreCache
+    dataStore: DataStore
     uuid: string
     actualWidth: number
     isProjectRun: boolean = false
@@ -46,6 +42,7 @@ class RunHeaderView extends ScreenView {
     noteField: EditableField
     sizeField: EditableField
     tagField: EditableField
+    dataStoreField: EditableField
     sizeCheckPoints: EditableField
     sizeTensorBoard: EditableField
     private deleteButton: DeleteButton
@@ -58,12 +55,14 @@ class RunHeaderView extends ScreenView {
         this.runCache = CACHE.getRun(this.uuid)
         this.statusCache = CACHE.getRunStatus(this.uuid)
         this.userCache = CACHE.getUser()
+        this.dataStoreCache = CACHE.getDataStore(this.uuid)
 
         this.deleteButton = new DeleteButton({onButtonClick: this.onDelete.bind(this), parent: this.constructor.name})
         this.loader = new DataLoader(async (force) => {
             this.status = await this.statusCache.get(force)
             this.run = await this.runCache.get(force)
             this.user = await this.userCache.get(force)
+            this.dataStore = await this.dataStoreCache.get(force)
         })
 
         this.editStatus = EditStatus.NOCHANGE
@@ -250,6 +249,15 @@ class RunHeaderView extends ScreenView {
                     name: 'Commit Message',
                     value: this.run.commit_message
                 }).render($)
+
+                this.dataStoreField = new EditableField({
+                    name: 'Data Store',
+                    value: JSON.stringify(this.dataStore.data),
+                    isEditable: true,
+                    onChange: this.onInputChange.bind(this),
+                    numEditRows: 10
+                })
+                this.dataStoreField.render($)
             })
         })
         this.deleteButton.hide(!(this.user.is_complete && this.run.is_claimed))
@@ -304,6 +312,23 @@ class RunHeaderView extends ScreenView {
             return
         }
 
+        let dataStore: any
+        try {
+            dataStore = JSON.parse(this.dataStoreField.getInput())
+        } catch (e) {
+            this.editStatus = EditStatus.CHANGE
+            UserMessages.shared.error("Data Store is not a valid JSON")
+            return
+        }
+
+        try {
+            this.dataStore = await this.dataStoreCache.update(dataStore)
+        } catch (e) {
+            this.editStatus = EditStatus.CHANGE
+            UserMessages.shared.networkError(e, "Failed to save data store")
+            return
+        }
+
         await this._render()
         this.editStatus = EditStatus.NOCHANGE
     }

app/ui/src/cache/cache.ts

@@ -8,9 +8,11 @@ import {SessionsList} from '../models/session_list'
 import {Session} from '../models/session'
 import {ProcessData} from "../analyses/sessions/process/types"
 import {Config} from "../models/config"
+import {DataStore} from "../models/data_store";
 
 const RELOAD_TIMEOUT = 60 * 1000
-const FORCE_RELOAD_TIMEOUT = -1
+const FORCE_RELOAD_TIMEOUT = -1 // Added to stop misuse from free users.
+// Ignored because there's no free version at the moment
 
 export function isReloadTimeout(lastUpdated: number): boolean {
     return (new Date()).getTime() - lastUpdated > RELOAD_TIMEOUT
@@ -108,6 +110,34 @@ export abstract class CacheObject<T> {
     }
 }
 
+export class DataStoreCache extends CacheObject<DataStore> {
+    private readonly runUUID: string
+
+    constructor(runUUID: string) {
+        super();
+        this.runUUID = runUUID
+    }
+    load(args: any): Promise<any> {
+        return NETWORK.getDataStore(args)
+    }
+
+    async get(isRefresh: boolean = false): Promise<DataStore> {
+        if (this.data == null || isRefresh) {
+            let res = await this.load(this.runUUID)
+            this.data = new DataStore(res)
+        }
+
+        return this.data
+    }
+
+    async update(data: any): Promise<DataStore> {
+        let res = await NETWORK.setDataStore(this.runUUID, data)
+        this.data = new DataStore(res)
+
+        return this.data
+    }
+}
+
 export class RunsListCache extends CacheObject<RunsList> {
     private loadedTags: Set<string>
 
@@ -589,6 +619,7 @@ class Cache {
     private sessions: { [uuid: string]: SessionCache }
     private runStatuses: { [uuid: string]: RunStatusCache }
     private sessionStatuses: { [uuid: string]: SessionStatusCache }
+    private dataStores: { [uuid: string]: DataStoreCache }
 
     private user: UserCache | null
     private runsList: RunsListCache | null
@@ -600,11 +631,20 @@ class Cache {
         this.runStatuses = {}
         this.sessionStatuses = {}
         this.customMetrics = {}
+        this.dataStores = {}
         this.runsList = null
         this.user = null
         this.sessionsList = null
     }
 
+    getDataStore(uuid: string) {
+        if (this.dataStores[uuid] == null) {
+            this.dataStores[uuid] = new DataStoreCache(uuid)
+        }
+
+        return this.dataStores[uuid]
+    }
+
     getCustomMetrics(uuid: string) {
         if (this.customMetrics[uuid] == null) {
             this.customMetrics[uuid] = new CustomMetricCache(uuid)

app/ui/src/models/data_store.ts

@@ -0,0 +1,7 @@
+export class DataStore {
+    data: Record<string, any>
+
+    constructor(data: Record<string, any>) {
+        this.data = data
+    }
+}

app/ui/src/network.ts

@@ -27,6 +27,14 @@ class Network {
         });
     }
 
+    async getDataStore(runUUID: string): Promise<any> {
+        return this.sendHttpRequest('GET', `/datastore/${runUUID}`)['promise']
+    }
+
+    async setDataStore(runUUID: string, data: object): Promise<any> {
+        return this.sendHttpRequest('POST', `/datastore/${runUUID}`, data)['promise']
+    }
+
     async getRun(runUUID: string): Promise<any> {
         return this.sendHttpRequest('GET', `/run/${runUUID}`)['promise']
     }

client/labml/app_api.py

@@ -211,3 +211,9 @@ def delete_custom_metric(self, run_uuid, metric_id):
 
     def get_logs(self, run_uuid: str, url: str, page_no: int):
         return self.network.send_http_request('POST', f'/logs/{url}/{run_uuid}', {'page': page_no})
+
+    def get_data_store(self, run_uuid):
+        return self.network.send_http_request('GET', f'/datastore/{run_uuid}')
+
+    def update_data_store(self, run_uuid, data):
+        return self.network.send_http_request('POST', f'/datastore/{run_uuid}', data)