Add logs for math verify pool

Author: SurbhiJainUSC

src/maxtext/trainers/post_train/rl/math_verify_pool.py

@@ -28,6 +28,8 @@
 import itertools
 import multiprocessing
 import os
+import uuid
+from etils import epath
 
 # Module-level persistent pool state.
 _POOL = None
@@ -71,7 +73,7 @@ def silent_worker_init():
     pass
 
 
-def verify_math_worker(golds, predictions):
+def verify_math_worker(idx, golds, predictions, parse_log_path, verify_log_path):
   """Worker-side math_verify grader."""
   try:
     from math_verify import parse, verify
@@ -80,6 +82,16 @@ def verify_math_worker(golds, predictions):
     gold_targets = (ExprExtractionConfig(), LatexExtractionConfig())
     pred_targets = (ExprExtractionConfig(), LatexExtractionConfig())
 
+    log_file = parse_log_path / f"{idx}.txt"
+    log_content = (
+        "START ============================\n"
+        f"Index: {idx}\n"
+        f"Golds: {golds}\n"
+        f"Predictions: {predictions}\n"
+        "END ==============================\n"
+    )
+    log_file.write_text(log_content)
+
     extracted_predictions = list(
         itertools.chain.from_iterable(parse(pred, pred_targets, parsing_timeout=None) for pred in predictions)
     )
@@ -89,7 +101,17 @@ def verify_math_worker(golds, predictions):
     if not extracted_predictions or not extracted_golds:
       return 0.0
 
-    print("Calling verify for extracted_golds: ", extracted_golds, " and predictions: ", extracted_predictions, " golds: ", golds, " predictions: ", predictions)
+    log_file = verify_log_path / f"{idx}.txt"
+    log_content = (
+        "START ============================\n"
+        f"Index: {idx}\n"
+        f"Golds: {golds}\n"
+        f"Extracted Golds: {extracted_golds}\n"
+        f"Predictions: {predictions}\n"
+        f"Extracted Predictions: {extracted_predictions}\n"
+        "END ==============================\n"
+    )
+    log_file.write_text(log_content)
     return max(
         (1.0 if any(verify(gold, pred, timeout_seconds=None) for gold in extracted_golds) else 0.0)
         for pred in extracted_predictions
@@ -150,6 +172,31 @@ def math_verify_pool(items, timeout=15, num_procs=None, log_fn=None):
 
   print(f"math_verify_pool called for {len(items)} items")
 
+  call_id = uuid.uuid4().hex
+  print("Call id: ", call_id)
+
+  logs_dir = "gs://nicogrande-maxtext-logs/debug/logs/run/rl/math_verify"
+  pool_log_path = epath.Path(tmvp_config.base_output_directory) / "pool"
+  pool_log_path.mkdir(parents=True, exist_ok=True)
+    
+  parse_log_path = epath.Path(tmvp_config.base_output_directory) / "parse" / f"{call_id}"
+  parse_log_path.mkdir(parents=True, exist_ok=True)
+
+  verify_log_path = epath.Path(tmvp_config.base_output_directory) / "verify" / f"{call_id}"
+  verify_log_path.mkdir(parents=True, exist_ok=True)
+
+  result_log_path = epath.Path(tmvp_config.base_output_directory) / "worker" / f"{call_id}"
+  result_log_path.mkdir(parents=True, exist_ok=True)
+
+  # log pool
+  log_file = pool_log_path / f"{call_id}.txt"
+  log_content = (
+      "START ============================\n"
+      f"Items: {items}\n"
+      "END ==============================\n"
+  )
+  log_file.write_text(log_content)
+
   cpu_count = multiprocessing.cpu_count()
   if num_procs is None:
     num_procs = min(_DEFAULT_MAX_PROCS, len(items), cpu_count)
@@ -161,10 +208,18 @@ def math_verify_pool(items, timeout=15, num_procs=None, log_fn=None):
   pool = _get_pool(num_procs)
   saw_timeout = False
   try:
-    jobs = [pool.apply_async(verify_math_worker, (golds, predictions)) for (_, golds, predictions) in items]
+    jobs = [pool.apply_async(verify_math_worker, (idx, golds, predictions, parse_log_path, verify_log_path)) for (idx, golds, predictions) in items]
     for i, job in enumerate(jobs):
       try:
         results[i] = float(job.get(timeout=timeout))
+        log_file = result_log_path / f"{i}.txt"
+        log_content = (
+            "START ============================\n"
+            f"Idx: {i}\n"
+            f"Score: {results[i]}\n"
+            "END ==============================\n"
+        )
+        log_file.write_text(log_content)
         cnt += 1
       except multiprocessing.TimeoutError:
         if log_fn is not None: