test flax is not reinstalled by notebooks

google · Feb 13, 2023 · bfaa805 · bfaa805
1 parent 26d46e4
commit bfaa805
Show file tree

Hide file tree

Showing 9 changed files with 57 additions and 25 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -72,7 +72,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install standalone dependencies only
       run: |
-        pip install .
+        pip install -e .
     - name: Test importing Flax
       run: |
         python -c "import flax"

diff --git a/docs/advanced_topics/contributing.md b/docs/advanced_topics/contributing.md
@@ -51,9 +51,12 @@ Follow these steps to contribute code:
    ```bash
    git clone https://github.com/YOUR_USERNAME/flax
    cd flax
-   pip install ".[testing]"
    pip install -e .
+   pip install ".[testing]"
    pip install -r docs/requirements.txt
+   # install in editable mode again because docs/requirements.txt
+   # reinstalls project in non-editable mode
+   pip install -e .
    ```
 
 5. Set up pre-commit hooks, this will run some automated checks during each `git` commit and
@@ -127,6 +130,10 @@ The former can be opened and executed directly in [Google Colab](https://colab.r
 Markdown makes it easier to track changes/diffs within version control and, for example, GitHub
 web UI, since `.ipynb` files are based on JSON.
 
+**NOTE**: If your notebook contains a cell that uses `pip` to install a package
+you must add a `skip-execution` tag to that cell so `myst-nb` will skip the cell
+when testing the notebooks.
+
 #### Editing Jupyter Notebooks (`.ipynb`)
 
 For making large changes that substantially modify code and outputs, it's recommended to edit

diff --git a/docs/guides/flax_on_pjit.ipynb b/docs/guides/flax_on_pjit.ipynb
@@ -43,12 +43,15 @@
    "cell_type": "code",
    "execution_count": 15,
    "metadata": {
-    "id": "867203db3bef"
+    "id": "867203db3bef",
+    "tags": [
+     "skip-execution"
+    ]
    },
    "outputs": [],
    "source": [
     "# Once Flax v0.6.4 is released, use `pip3 install flax`.\n",
-    "!pip3 install -qq \"git+https://github.com/google/flax.git@main#egg=flax\""
+    "! pip3 install -qq \"git+https://github.com/google/flax.git@main#egg=flax\""
    ]
   },
   {
@@ -975,7 +978,7 @@
    "formats": "ipynb,md:myst"
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -989,11 +992,11 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9 (main, Dec  7 2022, 13:47:07) [GCC 12.2.0]"
+   "version": "3.8.10"
   },
   "vscode": {
    "interpreter": {
-    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+    "hash": "ec7c69eb752b35b8fd728edc4753e382b54c10c43e6028c93b5837f81a552f5c"
    }
   }
  },

diff --git a/docs/guides/flax_on_pjit.md b/docs/guides/flax_on_pjit.md
@@ -7,7 +7,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.13.8
 kernelspec:
-  display_name: Python 3
+  display_name: .venv
   language: python
   name: python3
 ---
@@ -42,9 +42,10 @@ Install Flax from HEAD:
 
 ```{code-cell} ipython3
 :id: 867203db3bef
+:tags: [skip-execution]
 
 # Once Flax v0.6.4 is released, use `pip3 install flax`.
-!pip3 install -qq "git+https://github.com/google/flax.git@main#egg=flax"
+! pip3 install -qq "git+https://github.com/google/flax.git@main#egg=flax"
 ```
 
 +++ {"id": "a9601432b448"}

diff --git a/docs/guides/use_checkpointing.ipynb b/docs/guides/use_checkpointing.ipynb
@@ -35,13 +35,17 @@
    "cell_type": "code",
    "execution_count": 1,
    "id": "e80f8743",
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "skip-execution"
+    ]
+   },
    "outputs": [],
    "source": [
-    "!pip install -U -qq flax orbax\n",
+    "! pip install -U -qq flax orbax\n",
     "\n",
     "# Orbax needs to enable asyncio in a Colab environment.\n",
-    "!pip install -qq nest_asyncio"
+    "! pip install -qq nest_asyncio"
    ]
   },
   {
@@ -642,7 +646,7 @@
    "formats": "ipynb,md"
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -656,7 +660,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.15"
+   "version": "3.8.10"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "ec7c69eb752b35b8fd728edc4753e382b54c10c43e6028c93b5837f81a552f5c"
+   }
   }
  },
  "nbformat": 4,

diff --git a/docs/guides/use_checkpointing.md b/docs/guides/use_checkpointing.md
@@ -8,7 +8,7 @@ jupyter:
       format_version: '1.3'
       jupytext_version: 1.13.8
   kernelspec:
-    display_name: Python 3 (ipykernel)
+    display_name: .venv
     language: python
     name: python3
 ---
@@ -33,11 +33,11 @@ This guide covers the following:
 
 Install/upgrade Flax and [Orbax](https://github.com/google/orbax). For JAX installation with GPU/TPU support, visit [this section on GitHub](https://github.com/google/jax#installation).
 
-```python
-!pip install -U -qq flax orbax
+```python tags=["skip-execution"]
+! pip install -U -qq flax orbax
 
 # Orbax needs to enable asyncio in a Colab environment.
-!pip install -qq nest_asyncio
+! pip install -qq nest_asyncio
 ```
 
 Note: Before running `import jax`, create eight fake devices to mimic [multi-host environment](https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html?#aside-hosts-and-devices-in-jax) this notebook. Note that the order of imports is important here. The `os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'` command works only with the CPU backend. This means it won't work with GPU/TPU acceleration on if you're running this notebook in Google Colab. If you are already running the code on multiple devices (for example, in a 4x2 TPU environment), you can skip running the next cell.

diff --git a/docs/notebooks/flax_sharp_bits.ipynb b/docs/notebooks/flax_sharp_bits.ipynb
@@ -16,10 +16,14 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "tags": [
+     "skip-execution"
+    ]
+   },
    "outputs": [],
    "source": [
-    "!pip install -qq flax"
+    "! pip install -qq flax"
    ]
   },
   {
@@ -131,7 +135,7 @@
    "formats": "ipynb,md:myst"
   },
   "kernelspec": {
-   "display_name": "Python 3.10.4 64-bit",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -145,11 +149,11 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.8.10"
   },
   "vscode": {
    "interpreter": {
-    "hash": "3ad933181bd8a04b432d3370b9dc3b0662ad032c4dfaa4e4f1596c548f763858"
+    "hash": "ec7c69eb752b35b8fd728edc4753e382b54c10c43e6028c93b5837f81a552f5c"
    }
   }
  },

diff --git a/docs/notebooks/flax_sharp_bits.md b/docs/notebooks/flax_sharp_bits.md
@@ -7,7 +7,7 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.13.8
 kernelspec:
-  display_name: Python 3.10.4 64-bit
+  display_name: .venv
   language: python
   name: python3
 ---
@@ -21,7 +21,9 @@ Flax exposes the full power of JAX. And just like when using JAX, there are cert
 First, install and/or update Flax:
 
 ```{code-cell} ipython3
-!pip install -qq flax
+:tags: [skip-execution]
+
+! pip install -qq flax
 ```
 
 ## 🔪 `flax.linen.Dropout` layer and randomness

diff --git a/tests/run_all_tests.sh b/tests/run_all_tests.sh
@@ -85,6 +85,12 @@ if $RUN_DOCTEST; then
   pytest -n auto flax --doctest-modules --suppress-no-test-exit-code
 fi
 
+# check that flax is running on editable mode
+# (i.e. no notebook installed flax from pypi)
+echo "=== CHECKING FLAX IS EDITABLE ==="
+assert_error="flax is not running on editable mode."
+(cd docs; python -c "import flax; assert 'site-packages' not in flax.__file__, \"$assert_error\"")
+
 # env vars must be set after doctest
 export JAX_NUMPY_RANK_PROMOTION=raise
 export FLAX_PROFILE=1