From ae9341d47290abbd41510627db3f38ab02d45499 Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:11:48 -0700 Subject: [PATCH] Refresh exploration fields in monitor --- run_exploration_monitor.py | 64 ++++++++++++++++++++++++++- tests/test_run_exploration_monitor.py | 57 +++++++++++++++++++++++- 2 files changed, 119 insertions(+), 2 deletions(-) diff --git a/run_exploration_monitor.py b/run_exploration_monitor.py index 4842a5018d..7ee8c065c9 100644 --- a/run_exploration_monitor.py +++ b/run_exploration_monitor.py @@ -19,6 +19,7 @@ s - save current layout p - shows help menu I - view the associated exploration YAML file + R - refresh runs and add fields from the associated exploration YAML g - graphs first two rows L - graph & connect points sharing the 3rd column value 1–9 - graph & connect points sharing merged columns 3..(2+N) @@ -68,6 +69,48 @@ def load_runs(log_file: Path) -> List[Dict]: return docs +# Keys used by the exploration runner to compose groups, rather than arguments +# that are ultimately written into each run's ``config`` mapping. +EXPLORATION_SCHEMA_KEYS = { + "named_group", + "named_group_static", + "named_group_variations", + "named_group_alternates", +} + + +def load_exploration_fields(config_file: Path) -> set[str]: + """Return run-configuration field names declared by an exploration YAML. + + Exploration files support both a flat mapping and nested group syntax. A + field is therefore any mapping key whose value is not another mapping (or + a list of mappings), excluding the group-composition keys above. + """ + if not config_file.exists(): + return set() + + fields: set[str] = set() + + def visit(value) -> None: + if isinstance(value, dict): + for key, child in value.items(): + if isinstance(child, dict) or ( + isinstance(child, list) + and any(isinstance(item, dict) for item in child) + ): + visit(child) + elif key not in EXPLORATION_SCHEMA_KEYS: + fields.add(str(key)) + elif isinstance(value, list): + for item in value: + visit(item) + + with config_file.open() as f: + for document in yaml.safe_load_all(f): + visit(document) + return fields + + HOTKEYS_TEXT = ( "Enter: toggle sort by column\n" "h/l: move column left/right\n" @@ -87,6 +130,7 @@ def load_runs(log_file: Path) -> List[Dict]: "g: graph first two columns (opens a Plotly window)\n" "p: shows help menu\n" "I: view the associated exploration YAML file\n" + "R: refresh runs and add fields from the associated exploration YAML\n" "L: graph & connect points sharing the 3rd column value\n" "1–9: graph & connect points sharing merged columns 3..(2+N)\n" "q # #: multibarcharts - `q [1-9] [1-9]` - e.g. 'q 3 2' will create bar charts for columns 1 2 and 3, the next two columns (column 4 and column 5) as merged labels\n" @@ -505,7 +549,7 @@ def _compare_values(a, b) -> int: def refresh_table(self, new_cursor: Optional[int] = None) -> None: - """Reload data, apply sorting, and repopulate the DataTable.""" + """Reload data/config fields, apply sorting, and repopulate the table.""" if not self.table: return # Always reload the YAML log file so new runs appear @@ -513,6 +557,21 @@ def refresh_table(self, new_cursor: Optional[int] = None) -> None: if new_original != self.original_entries: self.original_entries = new_original + # Runs can gain config values over time, and the associated exploration + # can be edited before those runs have completed. Discover both sources + # on every refresh so new fields immediately become visible without + # discarding the user's existing column order or hidden-column choices. + discovered_keys = load_exploration_fields(self.exploration_config_file) + for entry in self.original_entries: + discovered_keys.update(entry.get("config", {}).keys()) + new_keys = sorted(discovered_keys.difference(self.all_columns)) + if new_keys: + self.param_keys = sorted(set(self.param_keys).union(new_keys)) + self.all_columns.extend(new_keys) + self.columns.extend( + col for col in new_keys if col not in self.hidden_cols + ) + # Re-apply any active row filters base_entries = list(self.original_entries) for col, op, val in self.row_filters: @@ -931,6 +990,9 @@ async def on_key(self, event: events.Key) -> None: self._msg(HOTKEYS_TEXT, timeout=10.0) elif key == "I": self.push_screen(ExplorationConfigScreen(self.exploration_config_file)) + elif key == "R": + self.refresh_table(new_cursor=c) + self._msg("Runs and exploration fields refreshed") elif key == "g": # ── Graph using first two visible columns: col[0] ⇒ Y, col[1] ⇒ X ── try: diff --git a/tests/test_run_exploration_monitor.py b/tests/test_run_exploration_monitor.py index 8c42442690..4adf52d0ea 100644 --- a/tests/test_run_exploration_monitor.py +++ b/tests/test_run_exploration_monitor.py @@ -1,7 +1,12 @@ +import tempfile import unittest from pathlib import Path -from run_exploration_monitor import ExplorationConfigScreen, MonitorApp +from run_exploration_monitor import ( + ExplorationConfigScreen, + MonitorApp, + load_exploration_fields, +) class ColumnSettingRemapTests(unittest.TestCase): @@ -45,5 +50,55 @@ async def test_hotkey_opens_associated_yaml_contents(self): self.assertIn("norm_variant_wte", str(contents)) +class ExplorationFieldRefreshTests(unittest.IsolatedAsyncioTestCase): + def test_loads_flat_and_grouped_fields_but_not_schema_keys(self): + with tempfile.TemporaryDirectory() as tmpdir: + config_file = Path(tmpdir) / "sweep.yaml" + config_file.write_text( + """\ +max_iters: [100] +named_static_groups: + - named_group: defaults + named_group_settings: + use_qk_norm: [true] +parameter_groups: + - learning_rate: [0.001] + named_group_static: [defaults] +""" + ) + + self.assertEqual( + load_exploration_fields(config_file), + {"max_iters", "use_qk_norm", "learning_rate"}, + ) + + async def test_refresh_and_hotkey_add_latest_exploration_fields(self): + with tempfile.TemporaryDirectory() as tmpdir: + root = Path(tmpdir) + log_file = root / "logs" / "sweep.yaml" + log_file.parent.mkdir() + log_file.write_text("config:\n max_iters: 100\n") + config_file = root / "explorations" / "sweep.yaml" + config_file.parent.mkdir() + config_file.write_text("max_iters: [100]\n") + + app = MonitorApp(log_file, interval=3600.0, csv_dir=tmpdir) + app.exploration_config_file = config_file + async with app.run_test() as pilot: + self.assertNotIn("learning_rate", app.columns) + + config_file.write_text( + "max_iters: [100]\nlearning_rate: [0.001]\n" + ) + app.refresh_table() + self.assertIn("learning_rate", app.columns) + + config_file.write_text( + "max_iters: [100]\nlearning_rate: [0.001]\ndropout: [0.1]\n" + ) + await pilot.press("R") + self.assertIn("dropout", app.columns) + + if __name__ == "__main__": unittest.main()