From 0644f39397f55fa16f0653a2d49f0b498a444859 Mon Sep 17 00:00:00 2001
From: Takahiro Ueda <t.ueda.od@juntendo.ac.jp>
Date: Tue, 13 Aug 2024 20:14:33 +0900
Subject: [PATCH] feat: set a timeout for each cell

---
 .github/workflows/run-tutorial-notebook.yml        | 7 ++++++-
 .github/workflows/scripts/run-tutorial-notebook.sh | 9 ++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/run-tutorial-notebook.yml b/.github/workflows/run-tutorial-notebook.yml
index cdd523c..6f58bf1 100644
--- a/.github/workflows/run-tutorial-notebook.yml
+++ b/.github/workflows/run-tutorial-notebook.yml
@@ -21,6 +21,11 @@ on:
         required: false
         default: ''
         type: string
+      timeout:
+        description: Timeout per cell in seconds
+        required: false
+        default: 3600
+        type: number
 
 jobs:
   run:
@@ -60,7 +65,7 @@ jobs:
         -v $(pwd)/madminer_shared:/home/shared
         -v $(pwd)/.github/workflows/scripts/run-tutorial-notebook.sh:/tmp/run-tutorial-notebook.sh:ro
         madminer-jupyter-env:latest
-        /tmp/run-tutorial-notebook.sh /home/shared/madminer/${{ inputs.working-directory }} ${{ inputs.notebook }}
+        /tmp/run-tutorial-notebook.sh /home/shared/madminer/${{ inputs.working-directory }} ${{ inputs.notebook }} ${{ inputs.timeout }}
 
     - name: Upload notebook and logs
       uses: actions/upload-artifact@v4
diff --git a/.github/workflows/scripts/run-tutorial-notebook.sh b/.github/workflows/scripts/run-tutorial-notebook.sh
index b0cc897..001fc14 100755
--- a/.github/workflows/scripts/run-tutorial-notebook.sh
+++ b/.github/workflows/scripts/run-tutorial-notebook.sh
@@ -4,6 +4,7 @@ set -eu
 
 working_directory=$1
 notebook_file=$2
+timeout=$3
 
 if [ ! -f /.dockerenv ]; then
   echo 'error: this script must be executed within a container' >&2
@@ -22,4 +23,10 @@ time {
 echo "::endgroup::"
 
 cd "$working_directory"
-papermill "$notebook_file" "$notebook_file"
+# We may set a timeout (for each cell) to ensure in-progress notebooks and logs are saved,
+# even if the workflow hits the total 6-hour limit in GitHub Actions.
+if [ "$timeout" -ne 0 ]; then
+  papermill --execution-timeout "$timeout" "$notebook_file" "$timeout"
+else
+  papermill "$notebook_file" "$timeout"
+fi