Merge pull request #1593 from Manik-Khajuria-5/Layer-2

Feat : Add experiment metadata model + store (Layer 2 foundation)

Author: hpdang

py/tests/test_experiment_store.py

@@ -0,0 +1,218 @@
+"""Unit tests for the experiments metadata model and store (Layer 2, PR 1).
+
+Exercises :mod:`visdom.experiments` against a real ``JSONStore`` over a
+temporary directory, so no running visdom server is needed. These cover the
+model schemas (validation and round-trip (de)serialisation) and the
+``ExperimentStore`` CRUD/list operations, including that experiment metadata
+persists to disk and coexists with ordinary environment window data.
+"""
+
+import tempfile
+import unittest
+
+from visdom.data_model import JSONStore
+from visdom.experiments import (
+    Experiment,
+    ExperimentStore,
+    Metric,
+    Param,
+    Tag,
+    STATUS_FAILED,
+    STATUS_FINISHED,
+    STATUS_RUNNING,
+)
+
+
+class TestModels(unittest.TestCase):
+    """The Experiment/Param/Metric/Tag schemas validate and round-trip."""
+
+    def test_param_infers_dtype(self):
+        """Param.dtype is inferred from the value when not given explicitly."""
+        self.assertEqual(Param("lr", 0.01).dtype, "float")
+        self.assertEqual(Param("epochs", 10).dtype, "int")
+        self.assertEqual(Param("name", "resnet").dtype, "str")
+        # bool must not be mislabelled as int (bool subclasses int).
+        self.assertEqual(Param("amp", True).dtype, "bool")
+
+    def test_param_respects_explicit_dtype(self):
+        """An explicit dtype is kept rather than being overwritten."""
+        self.assertEqual(Param("k", "1", dtype="int").dtype, "int")
+
+    def test_experiment_defaults(self):
+        """A new experiment defaults name to env_id and starts running."""
+        exp = Experiment(env_id="main")
+        self.assertEqual(exp.name, "main")
+        self.assertEqual(exp.status, STATUS_RUNNING)
+        self.assertIsNone(exp.finished_at)
+
+    def test_experiment_rejects_bad_status(self):
+        """Constructing with an unknown status raises ValueError."""
+        with self.assertRaises(ValueError):
+            Experiment(env_id="main", status="bogus")
+
+    def test_set_param_replaces_by_key(self):
+        """set_param overwrites an existing param of the same name in place."""
+        exp = Experiment(env_id="main")
+        exp.set_param("lr", 0.1)
+        exp.set_param("lr", 0.01)
+        self.assertEqual(len(exp.params), 1)
+        self.assertEqual(exp.get_param("lr").value, 0.01)
+
+    def test_set_tag_replaces_by_key(self):
+        """set_tag overwrites an existing tag of the same name in place."""
+        exp = Experiment(env_id="main")
+        exp.set_tag("dataset", "mnist")
+        exp.set_tag("dataset", "cifar")
+        self.assertEqual(len(exp.tags), 1)
+        self.assertEqual(exp.tags[0].value, "cifar")
+
+    def test_add_metric_appends_time_series(self):
+        """add_metric appends (never replaces); latest_metric reads the last."""
+        exp = Experiment(env_id="main")
+        exp.add_metric("acc", 0.5, step=1)
+        exp.add_metric("acc", 0.9, step=2)
+        self.assertEqual(len(exp.metrics), 2)
+        self.assertEqual(exp.latest_metric("acc").value, 0.9)
+        self.assertIsNone(exp.latest_metric("missing"))
+
+    def test_finish_sets_terminal_state(self):
+        """finish stamps finished_at and only accepts terminal statuses."""
+        exp = Experiment(env_id="main")
+        exp.finish()
+        self.assertEqual(exp.status, STATUS_FINISHED)
+        self.assertIsNotNone(exp.finished_at)
+        exp2 = Experiment(env_id="two")
+        exp2.finish(STATUS_FAILED)
+        self.assertEqual(exp2.status, STATUS_FAILED)
+        with self.assertRaises(ValueError):
+            Experiment(env_id="three").finish(STATUS_RUNNING)
+
+    def test_round_trip_serialisation(self):
+        """to_dict/from_dict preserve every field of a populated experiment."""
+        exp = Experiment(env_id="main", description="run 1")
+        exp.set_param("lr", 0.01)
+        exp.set_tag("dataset", "mnist")
+        exp.add_metric("acc", 0.9, step=3)
+        exp.finish()
+        rebuilt = Experiment.from_dict(exp.to_dict())
+        self.assertEqual(rebuilt.to_dict(), exp.to_dict())
+        self.assertIsInstance(rebuilt.params[0], Param)
+        self.assertIsInstance(rebuilt.metrics[0], Metric)
+        self.assertIsInstance(rebuilt.tags[0], Tag)
+
+
+class TestExperimentStore(unittest.TestCase):
+    """ExperimentStore CRUD/list operations over a real JSONStore backend."""
+
+    def setUp(self):
+        """Give each test a fresh temp env_path and a store over it."""
+        self._tmp = tempfile.TemporaryDirectory()
+        self.env_path = self._tmp.name
+        self.backend = JSONStore(self.env_path)
+        self.store = ExperimentStore(self.backend)
+
+    def tearDown(self):
+        self._tmp.cleanup()
+
+    def test_log_experiment_persists_to_disk(self):
+        """A logged experiment is readable by a brand-new store over the dir."""
+        self.store.log_experiment(
+            "main", params={"lr": 0.01, "epochs": 10}, tags={"dataset": "mnist"}
+        )
+        # A fresh store instance proves the data survives on disk, not just in
+        # the object we wrote through.
+        reopened = ExperimentStore(JSONStore(self.env_path))
+        exp = reopened.get_experiment("main")
+        self.assertIsNotNone(exp)
+        self.assertEqual(exp.get_param("lr").value, 0.01)
+        self.assertEqual(exp.get_param("epochs").dtype, "int")
+        self.assertEqual(exp.tags[0].value, "mnist")
+
+    def test_get_missing_experiment_returns_none(self):
+        """Envs with no experiment blob yield None (feature is opt-in)."""
+        self.assertIsNone(self.store.get_experiment("never_logged"))
+
+    def test_log_experiment_updates_in_place(self):
+        """Re-logging merges params and keeps prior metrics rather than resetting."""
+        self.store.log_experiment("main", params={"lr": 0.1})
+        self.store.log_metric("main", "acc", 0.7)
+        exp = self.store.log_experiment(
+            "main", description="updated", params={"lr": 0.01, "wd": 0.0}
+        )
+        self.assertEqual(exp.description, "updated")
+        self.assertEqual(exp.get_param("lr").value, 0.01)
+        self.assertEqual(exp.get_param("wd").value, 0.0)
+        # The metric logged before the update is still present.
+        self.assertEqual(len(exp.metrics), 1)
+
+    def test_log_metric_auto_creates_experiment(self):
+        """Logging a metric for an env with no experiment creates one."""
+        self.store.log_metric("main", "loss", 1.5, step=0)
+        exp = self.store.get_experiment("main")
+        self.assertIsNotNone(exp)
+        self.assertEqual(exp.latest_metric("loss").value, 1.5)
+
+    def test_finish_experiment(self):
+        """finish_experiment persists a terminal status."""
+        self.store.log_experiment("main")
+        self.store.finish_experiment("main", STATUS_FAILED)
+        self.assertEqual(self.store.get_experiment("main").status, STATUS_FAILED)
+
+    def test_finish_missing_experiment_raises(self):
+        """Finishing an env that never logged an experiment raises KeyError."""
+        with self.assertRaises(KeyError):
+            self.store.finish_experiment("nope")
+
+    def test_list_experiments(self):
+        """list_experiments returns only envs that actually have a blob."""
+        self.store.log_experiment("a")
+        self.store.log_experiment("b")
+        # An ordinary env with window data but no experiment must be skipped.
+        self.backend.save_env("plain", {"jsons": {"w": {"id": "w"}}, "reload": {}})
+        listed = sorted(exp.env_id for exp in self.store.list_experiments())
+        self.assertEqual(listed, ["a", "b"])
+
+    def test_delete_experiment_keeps_env(self):
+        """delete_experiment drops the blob but leaves the environment intact."""
+        self.backend.save_env("main", {"jsons": {"w": {"id": "w"}}, "reload": {}})
+        self.store.log_experiment("main", params={"lr": 0.01})
+        self.assertTrue(self.store.delete_experiment("main"))
+        self.assertIsNone(self.store.get_experiment("main"))
+        # The window data the env had before is untouched.
+        env = self.backend.load_env("main")
+        self.assertIn("w", env["jsons"])
+        self.assertNotIn("experiment", env)
+
+    def test_delete_missing_experiment_returns_false(self):
+        """Deleting from an env with no experiment reports False."""
+        self.assertFalse(self.store.delete_experiment("nope"))
+
+    def test_experiment_coexists_with_window_data(self):
+        """Logging onto an env with windows preserves those windows on disk."""
+        self.backend.save_env(
+            "main", {"jsons": {"win_0": {"id": "win_0"}}, "reload": {"foo": 1}}
+        )
+        self.store.log_experiment("main", params={"lr": 0.01})
+        env = self.backend.load_env("main")
+        self.assertIn("win_0", env["jsons"])
+        self.assertEqual(env["reload"], {"foo": 1})
+        self.assertEqual(env["experiment"]["params"][0]["value"], 0.01)
+
+
+class TestExperimentStoreNoPersistence(unittest.TestCase):
+    """With persistence disabled the store degrades gracefully (no crashes)."""
+
+    def setUp(self):
+        self.store = ExperimentStore(JSONStore(None))
+
+    def test_log_returns_experiment_but_read_is_empty(self):
+        """log_* still returns a valid object though nothing is persisted."""
+        exp = self.store.log_experiment("main", params={"lr": 0.01})
+        self.assertIsInstance(exp, Experiment)
+        # Nothing is stored, so a subsequent read finds no experiment.
+        self.assertIsNone(self.store.get_experiment("main"))
+        self.assertEqual(self.store.list_experiments(), [])
+
+
+if __name__ == "__main__":
+    unittest.main()

py/visdom/data_model/json_store.py

@@ -137,7 +137,12 @@ def save_all(self, state):
         return self.save_envs(state, list(state.keys()))
 
     def load_env(self, eid):
-        """Read one environment by ``eid``; return ``{}`` if it is absent."""
+        """Read one environment by ``eid``; return ``{}`` if it is absent.
+
+        Only the canonical env fields are returned: ``jsons`` and ``reload``,
+        plus the ``experiment`` metadata blob when present. Internal bookkeeping
+        such as the ``name`` field written for hashed long-id files is dropped.
+        """
         if self.env_path is None:
             return {}
         path = self._resolve_existing(eid)
@@ -148,9 +153,12 @@ def load_env(self, eid):
                 data = json.load(fn)
         except (OSError, ValueError):
             return {}
-        if isinstance(data, dict) and "jsons" in data and "reload" in data:
-            return {"jsons": data.get("jsons", {}), "reload": data.get("reload", {})}
-        return {}
+        if not (isinstance(data, dict) and "jsons" in data and "reload" in data):
+            return {}
+        env = {"jsons": data.get("jsons", {}), "reload": data.get("reload", {})}
+        if "experiment" in data:
+            env["experiment"] = data["experiment"]
+        return env
 
     def list_envs(self):
         """Return the ids of all environments stored on disk.

py/visdom/experiments/__init__.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+# Copyright 2017-present, The Visdom Authors
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from visdom.experiments.models import (
+    Experiment,
+    Metric,
+    Param,
+    Tag,
+    STATUS_FAILED,
+    STATUS_FINISHED,
+    STATUS_RUNNING,
+    VALID_STATUSES,
+)
+from visdom.experiments.store import ExperimentStore
+
+__all__ = [
+    "Experiment",
+    "ExperimentStore",
+    "Metric",
+    "Param",
+    "Tag",
+    "STATUS_FAILED",
+    "STATUS_FINISHED",
+    "STATUS_RUNNING",
+    "VALID_STATUSES",
+]