feat: add atomic backend storage and SDK for experiment tagging

js/api/ApiProvider.js

@@ -133,6 +133,12 @@ const ApiProvider = ({ children }) => {
       case 'env_update':
         apiHandlers.current.onEnvUpdate(cmd.data);
         break;
+      case 'tags_update':
+        apiHandlers.current.onTagsUpdate(cmd.data);
+        break;
+      case 'tags_sync':
+        apiHandlers.current.onTagsSync(cmd.data);
+        break;
 
       default:
         console.error('unrecognized command', cmd);

js/main.js

@@ -86,6 +86,7 @@ const App = () => {
   // data stores
   const [storeMeta, setStoreMeta] = useState({
     envList: ENV_LIST.slice(),
+    tags: typeof TAGS_INDEX !== 'undefined' ? TAGS_INDEX : {},
     layoutLists: new Map([['main', new Map([[DEFAULT_LAYOUT, new Map()]])]]),
   });
   const [storeData, setStoreData] = useState({
@@ -299,6 +300,22 @@ const App = () => {
       layoutLists: layoutLists,
     }));
   };
+  const onTagsUpdate = (data) => {
+    setStoreMeta((prev) => ({
+      ...prev,
+      tags: {
+        ...prev.tags,
+        [data.eid]: data.tags,
+      },
+    }));
+  };
+
+  const onTagsSync = (data) => {
+    setStoreMeta((prev) => ({
+      ...prev,
+      tags: data,
+    }));
+  };
 
   // remove paneID from pane list
   // (also tell server)
@@ -794,6 +811,7 @@ const App = () => {
     <EnvControls
       envIDs={selection.envIDs}
       envList={storeMeta.envList}
+      tags={storeMeta.tags}
       envSelectorStyle={{
         width: Math.max(window.innerWidth / 3, 50),
       }}
@@ -845,6 +863,8 @@ const App = () => {
     onLayoutMessage,
     onReloadMessage,
     onEnvUpdate,
+    onTagsUpdate,
+    onTagsSync,
     onCloseMessage,
     onDisconnect,
   };

js/topbar/EnvControls.js

@@ -18,6 +18,7 @@ function EnvControls(props) {
   const {
     envList,
     envIDs,
+    tags,
     envSelectorStyle,
     onEnvSelect,
     onEnvClear,
@@ -41,10 +42,14 @@ function EnvControls(props) {
     if (env.split('_').length == 1) {
       return null;
     }
+    const env_tags = (tags && tags[env]) || [];
+    const label = env_tags.length > 0 ? `${env} [${env_tags.join(', ')}]` : env;
+
     return {
       key: idx + 1 + roots.length,
       pId: roots.indexOf(env.split('_')[0]) + 1,
-      label: env,
+      label: label,
+      title: label,
       value: env,
     };
   });
@@ -53,10 +58,13 @@ function EnvControls(props) {
 
   env_options2 = env_options2.concat(
     roots.map((x, idx) => {
+      const root_tags = (tags && tags[x]) || [];
+      const label = root_tags.length > 0 ? `${x} [${root_tags.join(', ')}]` : x;
       return {
         key: idx + 1,
         pId: 0,
-        label: x,
+        label: label,
+        title: label,
         value: x,
       };
     })

py/visdom/__init__.py

@@ -815,6 +815,44 @@ def get_window_data(self, win=None, env=None):
             create=False,
         )
 
+    def set_tags(self, tags, env=None, append=False):
+        """
+        This function sets tags for a specified environment.
+        If append is True, tags are added to the existing ones.
+        Otherwise, tags are replaced.
+        """
+        if isinstance(tags, str):
+            tags = [tags]
+
+        return self._send(
+            msg={
+                "eid": env,
+                "tags": tags,
+                "append": append,
+            },
+            endpoint="tags",
+            create=False,
+        )
+
+    def get_tags(self, env=None):
+        """
+        This function returns the tags for a specified environment.
+        """
+        if env is None:
+            env = self.env
+
+        try:
+            url = "{0}:{1}{2}/tags?eid={3}".format(
+                self.server, self.port, self.base_url, env
+            )
+            r = self.session.get(url)
+            res = r.json()
+            # print(f"DEBUG SDK: get_tags type={type(res)} val={res}")
+            return res
+        except Exception as e:
+            # print(f"DEBUG SDK: get_tags error={e}")
+            return []
+
     def set_window_data(self, data, win=None, env=None):
         """
         This function sets all the window data for a specified window in

py/visdom/server/app.py

@@ -11,9 +11,11 @@
 all of the required state about the currently running server.
 """
 
+import json
 import logging
 import os
 import platform
+import threading
 import time
 
 import tornado.web  # noqa E402: gotta install ioloop first
@@ -42,6 +44,7 @@
     SaveHandler,
     UpdateHandler,
     UserSettingsHandler,
+    TagsHandler,
 )
 from visdom.server.defaults import (
     DEFAULT_BASE_URL,
@@ -74,6 +77,7 @@ def __init__(
     ):
         self.eager_data_loading = eager_data_loading
         self.env_path = env_path
+        self.index_lock = threading.Lock()
         self.state = self.load_state()
         self.layouts = self.load_layouts()
         self.user_settings = self.load_user_settings()
@@ -111,6 +115,7 @@ def __init__(
             (r"%s/delete_env" % self.base_url, DeleteEnvHandler, {"app": self}),
             (r"%s/env_state" % self.base_url, EnvStateHandler, {"app": self}),
             (r"%s/fork_env" % self.base_url, ForkEnvHandler, {"app": self}),
+            (r"%s/tags" % self.base_url, TagsHandler, {"app": self}),
             (r"%s/user/(.*)" % self.base_url, UserSettingsHandler, {"app": self}),
             (r"%s(.*)" % self.base_url, IndexHandler, {"app": self}),
         ]
@@ -163,7 +168,13 @@ def load_state(self):
             )
             return {"main": {"jsons": {}, "reload": {}}}
         ensure_dir_exists(env_path)
-        env_jsons = [i for i in os.listdir(env_path) if ".json" in i]
+        env_jsons = [
+            i
+            for i in os.listdir(env_path)
+            if i.endswith(".json") and i != "tags_index.json"
+        ]
+        self.tags = self.load_tag_index()
+
         for env_json in env_jsons:
             eid = env_json.replace(".json", "")
             env_path_file = os.path.join(env_path, env_json)
@@ -185,11 +196,33 @@ def load_state(self):
                 state[eid] = LazyEnvData(env_path_file)
 
         if "main" not in state and "main.json" not in env_jsons:
-            state["main"] = {"jsons": {}, "reload": {}}
+            state["main"] = {"jsons": {}, "reload": {}, "tags": []}
             serialize_env(state, ["main"], env_path=self.env_path)
 
         return state
 
+    def load_tag_index(self):
+        index_path = os.path.join(self.env_path, "tags_index.json")
+        if os.path.exists(index_path):
+            try:
+                with open(index_path, "r") as f:
+                    return json.load(f)
+            except Exception:
+                logging.warn(f"Failed to load tag index at {index_path}")
+        return {}
+
+    def save_tag_index(self):
+        index_path = os.path.join(self.env_path, "tags_index.json")
+        try:
+            from visdom.utils.server_utils import atomic_save
+
+            with self.index_lock:
+                atomic_save(index_path, json.dumps(self.tags))
+        except Exception:
+            import traceback
+
+            logging.warn(f"Failed to save tag index at {index_path}: {traceback.format_exc()}")
+
     def load_user_settings(self):
         settings = {}
 

py/visdom/server/handlers/socket_handlers.py

@@ -27,6 +27,7 @@
 from visdom.utils.server_utils import (
     check_auth,
     broadcast_envs,
+    sync_tags,
     serialize_env,
     send_to_sources,
     broadcast,
@@ -277,6 +278,7 @@ def open(self):
         )
         self.broadcast_layouts([self])
         broadcast_envs(self, [self])
+        sync_tags(self, [self])
 
     def broadcast_layouts(self, target_subs=None):
         if target_subs is None:

py/visdom/server/handlers/web_handlers.py

@@ -43,6 +43,8 @@
     compare_envs,
     load_env,
     broadcast,
+    broadcast_tags,
+    sync_tags,
     update_window,
     hash_password,
     stringify,
@@ -629,9 +631,11 @@ def initialize(self, app):
         self.user_credential = app.user_credential
         self.base_url = app.base_url if app.base_url != "" else "/"
         self.wrap_socket = app.wrap_socket
+        self.app = app
 
     def get(self, args, **kwargs):
         items = gather_envs(self.state, env_path=self.env_path)
+        tags_index = json.dumps(self.app.tags)
         if (not self.login_enabled) or self.current_user:
             """self.current_user is an authenticated user provided by Tornado,
             available when we set self.get_current_user in BaseHandler,
@@ -641,6 +645,7 @@ def get(self, args, **kwargs):
                 "index.html",
                 user=getpass.getuser(),
                 items=items,
+                tags_index=tags_index,
                 active_item="",
                 wrap_socket=self.wrap_socket,
             )
@@ -649,6 +654,7 @@ def get(self, args, **kwargs):
                 "login.html",
                 user=getpass.getuser(),
                 items=items,
+                tags_index=tags_index,
                 active_item="",
                 base_url=self.base_url,
             )
@@ -677,6 +683,65 @@ def get(self, path):
             self.write(self.user_settings["user_css"])
 
 
+class TagsHandler(BaseHandler):
+    def initialize(self, app):
+        self.state = app.state
+        self.env_path = app.env_path
+        self.subs = app.subs
+        self.login_enabled = app.login_enabled
+        self.app = app
+
+    @check_auth
+    def get(self):
+        """Handle GET requests for retrieving tags."""
+        eid = self.get_argument("eid", "main")
+        eid = escape_eid(eid)
+
+        if eid in self.state:
+            res = json.dumps(self.state[eid].get("tags", []))
+        elif eid in self.app.tags:
+            res = json.dumps(self.app.tags[eid])
+        else:
+            res = json.dumps([])
+
+        self.write(res)
+
+    @check_auth
+    def post(self):
+        """Handle POST requests for updating tags."""
+        args = tornado.escape.json_decode(
+            tornado.escape.to_basestring(self.request.body)
+        )
+        eid = extract_eid(args)
+        tags = args.get("tags", [])
+        append = args.get("append", False)
+
+        with self.app.index_lock:
+            if eid not in self.state:
+                self.state[eid] = {"jsons": {}, "reload": {}, "tags": []}
+
+            if append:
+                current_tags = self.state[eid].get("tags", [])
+                # Deduplicate while preserving order
+                new_tags = list(OrderedDict.fromkeys(current_tags + tags))
+                self.state[eid]["tags"] = new_tags
+            else:
+                self.state[eid]["tags"] = list(OrderedDict.fromkeys(tags))
+
+            # Update global index and save
+            self.app.tags[eid] = self.state[eid]["tags"]
+            self.app.save_tag_index()
+
+        # Broadcast update (outside the lock to minimize hold time)
+        broadcast_tags(self, eid, self.state[eid]["tags"])
+
+        # Async save env
+        serialize_env(self.state, [eid], env_path=self.env_path)
+
+        res = json.dumps(self.state[eid]["tags"])
+        self.write(res)
+
+
 class ErrorHandler(BaseHandler):
     def get(self, text):
         error_text = text or "test error"

py/visdom/static/index.html

@@ -51,6 +51,7 @@
       '{{escape(item)}}',
       {% end %}
     ];
+    var TAGS_INDEX = {% raw tags_index %};
     var ACTIVE_ENV = '{{escape(active_item)}}';
     var USER = '{{escape(user)}}';
     var USE_POLLING = ('{{wrap_socket}}' == 'True');