File saving to track end of evaluation runs. This helps with extensio…

…ns that might wait for some evals to finish before running an op.

PiperOrigin-RevId: 650927891

kauldron/evals/__init__.py

@@ -21,6 +21,9 @@
 # Lazy-import is import here as `run_strategies` is imported from kxm and
 # we do not want to trigger a full import.
 with _epy.lazy_api_imports(globals()):
+  from kauldron.evals.eval_impl import TRAIN_COMPLETE_FILENAME
+  from kauldron.evals.eval_impl import EVAL_COMPLETE_FILENAME
+
   from kauldron.evals.evaluators import CollectionKeys
   from kauldron.evals.evaluators import Evaluator
   from kauldron.evals.evaluators import EvaluatorBase

kauldron/evals/eval_impl.py

@@ -33,6 +33,7 @@
 # XManager API do not have API for jobs within a work-unit to communicate,
 # so use files for communication.
 TRAIN_COMPLETE_FILENAME = 'train_complete.txt'
+EVAL_COMPLETE_FILENAME = 'eval_{}_complete.txt'
 
 
 def continuous_eval(
@@ -120,11 +121,25 @@ def continuous_eval(
 
     final_step = step
 
+  # All every_checkpoint_evals have been processed. Marks those as complete.
+  if trainer.workdir.exists():  # `TrainEvaluator` do not have a workdir
+    for ev in every_checkpoint_evals:
+      epath.Path(trainer.workdir).joinpath(
+          EVAL_COMPLETE_FILENAME.format(ev.name)
+      ).touch()
+
   logging.info('Running final evals...')
   for ev in last_checkpoint_evals:
     with tracker.catch_exception(name=ev.name, step=final_step):
       aux[ev.name] = ev.evaluate(state=state, step=final_step)
 
+  # All last_checkpoint_evals have been processed. Marks those as complete.
+  if trainer.workdir.exists():  # `TrainEvaluator` do not have a workdir
+    for ev in last_checkpoint_evals:
+      epath.Path(trainer.workdir).joinpath(
+          EVAL_COMPLETE_FILENAME.format(ev.name)
+      ).touch()
+
   tracker.maybe_reraise()
 
   # Return the last aux