Add log output for jupyter notebooks

Author: Charles Li

.github/workflows/run_jupyter_notebooks.yml

@@ -73,7 +73,7 @@ jobs:
           # 2. Install MaxText package and all the post training dependencies
           uv pip install ${maxtext_wheel}[tpu-post-train] --resolution=lowest
           install_tpu_post_train_extra_deps
-          
+
           python3 -m pip freeze
       - name: Run Post-Training Notebooks
         shell: bash
@@ -115,7 +115,7 @@ jobs:
             echo "Running $filename ..."
             echo "------------------------------------------------------"
 
-            $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv
+            $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv --log-output --log-level INFO
           done
       - name: Upload Outputs
         if: always()

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -117,7 +117,7 @@
         "\n",
         "    # Install uv, a fast Python package installer\n",
         "    !pip install uv\n",
-        "    \n",
+        "\n",
         "    # Install MaxText and post-training dependencies\n",
         "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
         "    !install_tpu_post_train_extra_deps"
@@ -188,8 +188,10 @@
         "try:\n",
         "    from google.colab import userdata\n",
         "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
+        "    print(f\"get HF_TOKEN=f{HF_TOKEN} from google.colab\")\n",
         "except ImportError:\n",
         "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
+        "    print(f\"get HF_TOKEN=f{HF_TOKEN} from env\")\n",
         "\n",
         "# If not found in the environment, prompt the user for input securely\n",
         "# getpass function ensures the token is hidden while you type\n",