From 98c45aaed58e11fa09348f38a39dced48f5201b5 Mon Sep 17 00:00:00 2001
From: ashors1 <71393111+ashors1@users.noreply.github.com>
Date: Thu, 9 May 2024 09:18:34 -0700
Subject: [PATCH] Always remove eval checkpoints in test-pax.sh (#810)

Currently, eval checkpoints are only removed if the eval tests succeed.
In this PR, we remove the checkpoints in the case of failure as well.
---
 .github/container/test-pax.sh | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/container/test-pax.sh b/.github/container/test-pax.sh
index 5562b729a..91cb926b4 100755
--- a/.github/container/test-pax.sh
+++ b/.github/container/test-pax.sh
@@ -388,6 +388,9 @@ else
 fi
 
 if [[ ${EVALUATE} -ne 0 ]]; then
+
+  trap "rm -rf ${OUTPUT}/checkpoints" ERR INT HUP TERM EXIT
+
   ## train for 0 steps to generate an initial checkpoint
   python -m paxml.main \
     --fdl_config=${CONFIG} \
@@ -408,7 +411,6 @@ if [[ ${EVALUATE} -ne 0 ]]; then
     $ADDITIONAL_ARGS \
     $([[ $MULTIPROCESS != 0 ]] && echo --multiprocess_gpu)
 
-  rm -rf ${OUTPUT}/checkpoints
 else
   python -m paxml.main \
     --fdl_config=${CONFIG} \