add jsonl log reloading for continued runs

Author: TablewareBox

evals/eval.py

@@ -27,14 +27,14 @@
 _MAX_SAMPLES = None
 
 
-def _index_samples(samples: List[Any]) -> List[Tuple[Any, int]]:
+def _index_samples(samples: List[Any], completed_ids: List[int] = []) -> List[Tuple[Any, int]]:
     """Shuffle `samples` and pair each sample with its index."""
     indices = list(range(len(samples)))
     random.Random(SHUFFLE_SEED).shuffle(indices)
     if _MAX_SAMPLES is not None:
         indices = indices[:_MAX_SAMPLES]
     logger.info(f"Evaluating {len(indices)} samples")
-    work_items = [(samples[i], i) for i in indices]
+    work_items = [(samples[i], i) for i in indices if not i in completed_ids]
     return work_items
 
 
@@ -120,7 +120,10 @@ def eval_all_samples(
         """
         Evaluate all provided samples in parallel.
         """
-        work_items = _index_samples(samples)
+        samples_completed_ids = [int(event.sample_id.split(".")[-1]) for event in recorder.get_events(type="sampling")]
+        logger.info(f"Completed samples: {samples_completed_ids}")
+
+        work_items = _index_samples(samples, completed_ids=samples_completed_ids)
         threads = int(os.environ.get("EVALS_THREADS", "10"))
         show_progress = bool(os.environ.get("EVALS_SHOW_EVAL_PROGRESS", show_progress))
 
@@ -207,7 +210,10 @@ def eval_all_samples(
         """
         Evaluate all provided samples in parallel.
         """
-        work_items = _index_samples(samples)
+        samples_completed_ids = [int(event.sample_id.split(".")[-1]) for event in recorder.get_events(type="sampling")]
+        logger.info(f"Completed samples: {samples_completed_ids}")
+
+        work_items = _index_samples(samples, completed_ids=samples_completed_ids)
         threads = int(os.environ.get("EVALS_THREADS", "10"))
         show_progress = bool(os.environ.get("EVALS_SHOW_EVAL_PROGRESS", show_progress))
 

evals/record.py

@@ -9,7 +9,9 @@
 import atexit
 import contextlib
 import dataclasses
+import json
 import logging
+import os.path
 import threading
 import time
 from contextvars import ContextVar
@@ -340,8 +342,13 @@ def __init__(
         self.event_file_path = log_path
         self.hidden_data_fields = hidden_data_fields
         if log_path is not None:
-            with bf.BlobFile(log_path, "wb") as f:
-                f.write((jsondumps({"spec": dataclasses.asdict(run_spec)}) + "\n").encode("utf-8"))
+            if not os.path.exists(log_path):
+                with bf.BlobFile(log_path, "wb") as f:
+                    f.write((jsondumps({"spec": dataclasses.asdict(run_spec)}) + "\n").encode("utf-8"))
+            else:
+                lines = bf.BlobFile(log_path, "rb").readlines()
+                run_spec.run_id = json.loads(lines[0])["spec"]["run_id"]
+                self._events = [Event(**json.loads(line)) for line in lines[1:]]
 
     def _flush_events_internal(self, events_to_write: Sequence[Event]):
         start = time.time()