test the upload runs

.github/workflows/notebook-run.yml

@@ -29,6 +29,9 @@ env:
 
 jobs:
   notebook-run:
+    # Due to github worker hard limitation, of 24 hours
+    # we apply a timeout of 23 hours instead.
+    timeout-minutes: 1380
     name: ${{github.event.inputs.notebookFile}} on ${{github.event.inputs.gpuTarget}} / cuda-${{github.event.inputs.cudaVersion}}
     runs-on: 
       - cuda-${{github.event.inputs.cudaVersion}}

notebook/github-runner/github-runner.sh

@@ -14,6 +14,11 @@ if [[ -z "${WANDB_API_KEY}" ]]; then
     exit 1
 fi
 
+# The HF repo directory to use
+if [[ -z "${HF_REPO_SYNC}" ]]; then
+    HF_REPO_SYNC="rwkv-x-dev/rwkv-x-playground"
+fi
+
 # Get the notebook script from the first arg
 NOTEBOOK_FILE=$1
 
@@ -31,15 +36,15 @@ CACHE_DIR="$ACTION_DIR/.cache/"
 mkdir -p "$CACHE_DIR"
 
 # Log the proj dir
-echo "#"
+echo "# ------"
 echo "# Starting github notebook runner"
 echo "#"
 echo "# PROJ_DIR: $PROJ_DIR"
 echo "# NOTEBOOK_DIR: $NOTEBOOK_DIR"
 echo "# NOTEBOOK_FILE: $NOTEBOOK_FILE"
 echo "#"
 echo "# CACHE_DIR: $CACHE_DIR"
-echo "#"
+echo "# ------"
 
 # Check if the notebook file exists, in the notebook directory
 if [[ ! -f "$NOTEBOOK_DIR/$NOTEBOOK_FILE" ]]; then
@@ -127,4 +132,18 @@ echo "# [NOTE] Running notebook: $NOTEBOOK_FILE"
 cd "$INPUT_FILE_DIR"
 papermill \
     -k python3 --log-output \
-    "$INPUT_FILE_PATH" "$OUTPUT_FILE_PATH" 
+    "$INPUT_FILE_PATH" "$OUTPUT_FILE_PATH" 
+
+# -----
+# Upload the output notebook to the github repo
+# -----
+
+# Upload the result files
+echo "# ------"
+echo "# Uploading models & notebooks to HF repo"
+echo "# ------"
+
+# Get $NOTEBOOK_FILE, without the ipynb filetype
+NOTEBOOK_FILE_NOEXT="${NOTEBOOK_FILE%.*}"
+
+python3 ./hf-upload.py "$HF_REPO_SYNC" "$NOTEBOOK_FILE_NOEXT"

notebook/github-runner/hf-upload.py

@@ -0,0 +1,40 @@
+# Get the Hugging Face Hub API
+from huggingface_hub import HfApi
+api = HfApi()
+
+# Get the repo path from the script first arg
+import sys
+REPO_PATH = sys.argv[1]
+REPO_SUBDIR = sys.argv[2]
+
+# Get the current script dir
+import os
+
+RUNNER_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+NOTEBOOK_DIR = os.path.dirname(RUNNER_SCRIPT_DIR)
+PROJ_DIR = os.path.dirname(NOTEBOOK_DIR)
+
+MODEL_DIR = os.path.join(PROJ_DIR, "model")
+OUTPUT_DIR = os.path.join(PROJ_DIR, "output")
+
+# Upload the models
+api.upload_folder(
+    folder_path=MODEL_DIR,
+    repo_id=REPO_PATH,
+    path_in_repo=REPO_SUBDIR,
+    repo_type="model",
+    multi_commits=True,
+    allow_patterns=["*.pth"],
+    commit_message=f"[GHA] {REPO_SUBDIR}.ipynb result models"
+)
+
+# Upload the ipynb files
+api.upload_folder(
+    folder_path=OUTPUT_DIR,
+    repo_id=REPO_PATH,
+    path_in_repo=REPO_SUBDIR,
+    repo_type="model",
+    multi_commits=True,
+    allow_patterns=["*.ipynb"],
+    commit_message=f"[GHA] {REPO_SUBDIR}.ipynb result notebooks"
+)