Add jupyter notebook tests to Github CI tests

Author: SurbhiJainUSC

.github/workflows/build_and_test_maxtext.yml

@@ -67,6 +67,7 @@ jobs:
             echo "No files changed"
             echo "run_tests=true" >> $GITHUB_OUTPUT
             echo "run_notebooks=true" >> $GITHUB_OUTPUT
+            exit 0
           fi
 
           # Check for source code changes (anything not .md and not .ipynb)
@@ -87,6 +88,8 @@ jobs:
             echo "run_notebooks=false" >> $GITHUB_OUTPUT
           fi
 
+          exit 0
+
   build_and_upload_maxtext_package:
     needs: doc_only_check
     # Run if either tests or notebooks need to run

src/MaxText/examples/rl_llama3_demo.ipynb

@@ -363,7 +363,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

src/MaxText/examples/sft_llama3_demo.ipynb

@@ -173,7 +173,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
    "source": [
     "if IN_COLAB:\n",
     "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
@@ -345,7 +344,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "base",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },

src/MaxText/examples/sft_qwen3_demo.ipynb

@@ -111,6 +111,21 @@
     "**⚠️ Note:** The installation process in following cell may take a few minutes to complete. Please be patient."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "try:\n",
+    "    import google.colab\n",
+    "    print(\"Running the notebook in Google Colab\")\n",
+    "    IN_COLAB = True\n",
+    "except ImportError:\n",
+    "    print(\"Running the notebook on JupyterLab\")\n",
+    "    IN_COLAB = False"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -586,7 +601,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },