From 246c7641dafcc1dfbb2c22dc92099d53086b3502 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 28 Jul 2026 15:41:41 -0700 Subject: [PATCH] fix(dashboard): isolate telemetry per training run --- miles/dashboard/backend.py | 6 +- miles/dashboard/collector.py | 4 ++ miles/dashboard/hooks.py | 2 +- miles/dashboard/store.py | 15 +++++ miles/dashboard/viewer.py | 8 ++- miles/utils/arguments.py | 2 +- tests/fast/dashboard/test_async_collector.py | 1 - tests/fast/dashboard/test_run_isolation.py | 71 ++++++++++++++++++++ 8 files changed, 101 insertions(+), 8 deletions(-) create mode 100644 tests/fast/dashboard/test_run_isolation.py diff --git a/miles/dashboard/backend.py b/miles/dashboard/backend.py index 9417d789..289c07fe 100644 --- a/miles/dashboard/backend.py +++ b/miles/dashboard/backend.py @@ -6,6 +6,7 @@ import time from miles.dashboard.collector import COLLECTOR_ACTOR_NAME, CollectorConfig +from miles.dashboard.store import run_dir logger = logging.getLogger(__name__) @@ -27,10 +28,11 @@ def init_dashboard(args) -> bool: from miles.dashboard.collector import DashboardCollector + start_ts = time.time() config = CollectorConfig( - workspace=args.miles_dashboard_workspace, + workspace=str(run_dir(args.miles_dashboard_workspace, start_ts)), run_name="miles-diffusion", - start_ts=time.time(), + start_ts=start_ts, args_snapshot=dict(vars(args)), ) handle = None diff --git a/miles/dashboard/collector.py b/miles/dashboard/collector.py index 498c123f..245ba30e 100644 --- a/miles/dashboard/collector.py +++ b/miles/dashboard/collector.py @@ -41,6 +41,10 @@ def __init__(self, config: CollectorConfig) -> None: def ping(self) -> bool: return True + def workspace(self) -> str: + """This run's directory, for workers that write telemetry files themselves.""" + return self.config.workspace + def start(self) -> None: if self._flush_thread is not None: return diff --git a/miles/dashboard/hooks.py b/miles/dashboard/hooks.py index d85d8dd9..e4566d65 100644 --- a/miles/dashboard/hooks.py +++ b/miles/dashboard/hooks.py @@ -138,7 +138,7 @@ def register_train_actor(args, role: str) -> None: attach_phase_sink(handle, role) global _GPU_SAMPLER if dist.get_rank() == 0 and _GPU_SAMPLER is None: - _GPU_SAMPLER = GpuUtilSampler(args.miles_dashboard_workspace) + _GPU_SAMPLER = GpuUtilSampler(_ray_get(handle.workspace.remote())) _GPU_SAMPLER.start() diff --git a/miles/dashboard/store.py b/miles/dashboard/store.py index d1624b18..c08dc869 100644 --- a/miles/dashboard/store.py +++ b/miles/dashboard/store.py @@ -11,6 +11,9 @@ from miles.dashboard.events import PhaseEvent, TrajectoryEvent +RUN_DIR_PREFIX = "run_" + + class Stream(StrEnum): PHASES = "phases" TRAJECTORIES = "trajectories" @@ -35,6 +38,18 @@ def _hour_key(timestamp: float) -> str: return datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y%m%d_%H") +def run_dir(base: str, start_ts: float) -> Path: + """One directory per training launch, under the base workspace.""" + stamp = datetime.fromtimestamp(start_ts, tz=timezone.utc).strftime("%Y%m%d_%H%M%S") + return Path(base) / f"{RUN_DIR_PREFIX}{stamp}" + + +def resolve_run_dir(base: str) -> Path: + """Newest run under the base workspace (names sort by time), or base itself if it holds no runs.""" + runs = sorted(path for path in Path(base).glob(f"{RUN_DIR_PREFIX}*") if path.is_dir()) + return runs[-1] if runs else Path(base) + + class DashboardStore: def __init__(self, workspace: str) -> None: self.workspace = Path(workspace) diff --git a/miles/dashboard/viewer.py b/miles/dashboard/viewer.py index 93c070ab..ca9384b1 100644 --- a/miles/dashboard/viewer.py +++ b/miles/dashboard/viewer.py @@ -12,6 +12,7 @@ from collections import defaultdict from miles.dashboard.events import SPAN_KINDS +from miles.dashboard.store import resolve_run_dir def _read_jsonl(paths): @@ -324,7 +325,7 @@ def log_message(self, *a): def main(): ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--workspace", required=True, help="miles dashboard workspace") + ap.add_argument("--workspace", required=True, help="miles dashboard workspace, or one run directory in it") ap.add_argument("--out", default="dashboard.html") ap.add_argument("--title", default="miles-D rollout dashboard") ap.add_argument( @@ -334,10 +335,11 @@ def main(): ap.add_argument("--port", type=int, default=8000) args = ap.parse_args() + workspace = str(resolve_run_dir(args.workspace)) if args.serve: - serve(args.workspace, args.host, args.port, args.title) + serve(workspace, args.host, args.port, args.title) return - phases, gpu, life = load_streams(args.workspace) + phases, gpu, life = load_streams(workspace) html = build_html(phases, gpu, life, title=args.title) with open(args.out, "w") as f: f.write(html) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index e3a583d1..63672b41 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1105,7 +1105,7 @@ def add_debug_arguments(parser): "--miles-dashboard-workspace", type=str, default="./miles_dashboard", - help="Directory for miles dashboard telemetry.", + help="Base directory for miles dashboard telemetry; each launch writes to its own run_ subdir.", ) return parser diff --git a/tests/fast/dashboard/test_async_collector.py b/tests/fast/dashboard/test_async_collector.py index c136a9da..54d15047 100644 --- a/tests/fast/dashboard/test_async_collector.py +++ b/tests/fast/dashboard/test_async_collector.py @@ -80,7 +80,6 @@ def remote(self, config): monkeypatch.setitem(sys.modules, "ray", ray_module) monkeypatch.setitem(sys.modules, "ray.util.scheduling_strategies", strategies_module) monkeypatch.setattr(backend, "_handle", None) - monkeypatch.setattr(backend, "_is_primary", False) args = SimpleNamespace( use_miles_dashboard=True, diff --git a/tests/fast/dashboard/test_run_isolation.py b/tests/fast/dashboard/test_run_isolation.py new file mode 100644 index 00000000..9203018c --- /dev/null +++ b/tests/fast/dashboard/test_run_isolation.py @@ -0,0 +1,71 @@ +import sys +import types +from types import SimpleNamespace + +from miles.dashboard import backend +from miles.dashboard.collector import CollectorConfig, DashboardCollector +from miles.dashboard.events import PhaseEvent +from miles.dashboard.store import resolve_run_dir, run_dir +from miles.dashboard.viewer import load_streams + + +def _write_run(base, start_ts: float, phase_name: str): + workspace = run_dir(str(base), start_ts) + collector = DashboardCollector(CollectorConfig(workspace=str(workspace), run_name="test", start_ts=start_ts)) + collector.push_phases( + [PhaseEvent(name=phase_name, t0=start_ts, t1=start_ts + 1, role="actor", pid=10, node="n", gpus=[2], rank=7)] + ) + collector.flush() + return workspace + + +def test_relaunch_does_not_merge_metrics(tmp_path): + first = _write_run(tmp_path, 1_800_000_000.0, "first_run_phase") + second = _write_run(tmp_path, 1_800_000_060.0, "second_run_phase") + + assert first != second + assert [phase["name"] for phase in load_streams(str(first))[0]] == ["first_run_phase"] + assert [phase["name"] for phase in load_streams(str(second))[0]] == ["second_run_phase"] + + +def test_resolve_run_dir_picks_newest_run(tmp_path): + _write_run(tmp_path, 1_800_000_000.0, "old") + newest = _write_run(tmp_path, 1_800_000_060.0, "new") + + assert resolve_run_dir(str(tmp_path)) == newest + assert resolve_run_dir(str(newest)) == newest + + flat = tmp_path / "no_runs_inside" + (flat / "phases").mkdir(parents=True) + assert resolve_run_dir(str(flat)) == flat + + +def test_init_dashboard_writes_under_the_given_workspace(tmp_path, monkeypatch): + created = [] + + class FakeHandle: + ping = start = SimpleNamespace(remote=lambda *a, **kw: None) + + class RemoteClass: + def options(self, **kwargs): + return self + + def remote(self, config): + created.append(config) + return FakeHandle() + + ray_module = types.ModuleType("ray") + ray_module.remote = lambda cls: RemoteClass() + ray_module.get_runtime_context = lambda: SimpleNamespace(get_node_id=lambda: "driver-node") + ray_module.get = lambda ref: ref + strategies_module = types.ModuleType("ray.util.scheduling_strategies") + strategies_module.NodeAffinitySchedulingStrategy = lambda **kwargs: kwargs + monkeypatch.setitem(sys.modules, "ray", ray_module) + monkeypatch.setitem(sys.modules, "ray.util.scheduling_strategies", strategies_module) + monkeypatch.setattr(backend, "_handle", None) + + args = SimpleNamespace(use_miles_dashboard=True, miles_dashboard_workspace=str(tmp_path)) + assert backend.init_dashboard(args) is True + + assert created[0].workspace == str(run_dir(str(tmp_path), created[0].start_ts)) + assert args.miles_dashboard_workspace == str(tmp_path)