Remove GCS logging

Author: SurbhiJainUSC

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

@@ -94,7 +94,7 @@ def are_equal_under_sympy(gold, prediction):
     return False
 
 
-def verify_math_worker(idx, golds, predictions, parse_log_path=None, verify_log_path=None):
+def verify_math_worker(idx, golds, predictions):
   """Worker-side math_verify grader."""
   try:
     from math_verify import parse, verify
@@ -104,17 +104,6 @@ def verify_math_worker(idx, golds, predictions, parse_log_path=None, verify_log_
     gold_targets = (ExprExtractionConfig(), LatexExtractionConfig())
     pred_targets = (ExprExtractionConfig(), LatexExtractionConfig())
 
-    if parse_log_path is not None:
-      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)
     )
@@ -124,19 +113,6 @@ def verify_math_worker(idx, golds, predictions, parse_log_path=None, verify_log_
     if not extracted_predictions or not extracted_golds:
       return idx, 0.0
 
-    if verify_log_path is not None:
-      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)
-
     for gold in extracted_golds:
       for pred in extracted_predictions:
         if isinstance(gold, (Basic, MatrixBase)) and isinstance(pred, (Basic, MatrixBase)):
@@ -228,29 +204,6 @@ def math_verify_pool(tmvp_config, items, scores, timeout=15, num_procs=None, log
     return scores
 
   call_id = uuid.uuid4().hex
-  logs_dir = tmvp_config.base_output_directory + tmvp_config.run_name + "/math_verify"
-  pool_log_path = epath.Path(logs_dir) / "pool"
-  pool_log_path.mkdir(parents=True, exist_ok=True)
-    
-  parse_log_path = epath.Path(logs_dir) / "parse" / f"{call_id}"
-  parse_log_path.mkdir(parents=True, exist_ok=True)
-
-  verify_log_path = epath.Path(logs_dir) / "verify" / f"{call_id}"
-  verify_log_path.mkdir(parents=True, exist_ok=True)
-
-  result_log_path = epath.Path(logs_dir) / "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"Total Items: {len(items)}\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)
@@ -260,7 +213,7 @@ def math_verify_pool(tmvp_config, items, scores, timeout=15, num_procs=None, log
   cnt = 0
   timout_job_cnt = 0
   pool = _get_pool(num_procs)
-  active_jobs = [(idx, pool.apply_async(verify_math_worker, (idx, golds, predictions, parse_log_path, verify_log_path))) for (idx, golds, predictions) in items]
+  active_jobs = [(idx, pool.apply_async(verify_math_worker, (idx, golds, predictions))) for (idx, golds, predictions) in items]
   start_time = time.time()
   global_timeout = 300
   while active_jobs and (time.time() - start_time < global_timeout):
@@ -271,14 +224,6 @@ def math_verify_pool(tmvp_config, items, scores, timeout=15, num_procs=None, log
           # .get(0) returns immediately since ready() was true
           _, score = job.get(0)
           scores[idx] = max(scores[idx], tmvp_config.reward_exact_answer)
-          log_file = result_log_path / f"{idx}.txt"
-          log_content = (
-              "START ============================\n"
-              f"Idx: {idx}\n"
-              f"Score: {scores[idx]}\n"
-              "END ==============================\n"
-          )
-          log_file.write_text(log_content)
           cnt += 1
         except Exception as e:
           if log_fn is not None: