1 files changed, 7 insertions, 3 deletions
diff --git a/ep_run/sync_wandb.py b/ep_run/sync_wandb.py
index 74fff35..fc18bbc 100644
--- a/ep_run/sync_wandb.py
+++ b/ep_run/sync_wandb.py
@@ -9,7 +9,7 @@ ENTITY = 'eqprop-llm-training'
 R = Path('/home/yurenh2/ept/ep_run/runs')
 EP_RE = re.compile(
     r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) '
-    r'\| beta=([\d.e+-]+) K=(\d+) skips=(\d+)\(d(\d+)/g(\d+)/r(\d+)\)(?: cos=([\d.]+))? '
+    r'\| beta=([\d.e+-]+) K=(\d+) skips=(\d+)\(d(\d+)/g(\d+)/r(\d+)\)(?: cos=(-?[\d.]+))? '
     r'drift=([\d.]+) gn=([\d.e+-]+) sig=([\d.]+) \| ([\d.]+) it/s')
 BP_RE = re.compile(r'^step\s+(\d+)/\d+ \| train ([\d.]+) val ([\d.]+) \(best ([\d.]+)\) \| ([\d.]+) it/s')
 HDR = re.compile(r'^\[(\w+)\] cascade-')
@@ -34,7 +34,8 @@ def sync_file(path, default_tag, project, follow=False):
     def ensure(t):
         nonlocal run, last
         if run is not None: run.finish()
-        run = wandb.init(entity=ENTITY, project=project, name=t, id=t, resume='allow', reinit=True)
+        import os as _os
+        run = wandb.init(entity=ENTITY, project=project, name=t, id=t + _os.environ.get('SYNC_ID_SUFFIX', ''), resume='allow', reinit=True)
         last = -1
     f = open(path, errors='replace')
     ensure(tag)
@@ -59,7 +60,7 @@ HISTORY = [
     ('amp_s1.log', 'amp_s1', TS), ('amp_s2.log', 'amp_s2', TS), ('amp_s3.log', 'amp_s3', TS),
     ('nsize_ladder.log', 'nsize_c256', TS),
     ('bsign_s1.log', 'bsign_s1', TS), ('bsign_s2.log', 'bsign_s2', TS), ('bsign_s3.log', 'bsign_s3', TS),
-    ('f3e3cont.log', 'f3e3cont', TS),
+    ('stage1b_f3e3cont.log', 'f3e3cont', TS),
     ('arm_ctl.log', 'arm_ctl', TS), ('arm_mom99_f3e4.log', 'arm_mom99_f3e4', TS),
     ('arm_mom995_f1e4.log', 'arm_mom995_f1e4', TS), ('arm_cent_f3e3.log', 'arm_cent_f3e3', TS),
     ('arm_rich_f1e3.log', 'arm_rich_f1e3', TS), ('arm_mom99_f1e3.log', 'arm_mom99_f1e3', TS),
@@ -70,6 +71,9 @@ HISTORY = [
     ('w2_ctl.log', 'w2_ctl', FW), ('w2_blow.log', 'w2_blow', FW), ('w2_geta07.log', 'w2_geta07', FW),
     ('w2_geta05_k5.log', 'w2_geta05_k5', FW), ('w2_bcap.log', 'w2_bcap', FW), ('w2_adapt.log', 'w2_adapt', FW),
 ]
+import sys
+if len(sys.argv) > 1:
+    HISTORY = [h for h in HISTORY if h[1] in sys.argv[1:]]
 for fn, tag, prj in HISTORY:
     p = R / fn
     if p.exists():