From 593c2f1d453ec7cc51fdce75999782175f04fd97 Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Fri, 19 Apr 2024 13:19:45 +0200
Subject: [PATCH] ci/docs: set Sphinx public theme as alternative (#236)

---
 .github/workflows/docs-build.yml |  9 ++-------
 Makefile                         |  4 ++--
 docs/source/conf.py              | 25 +++++++++++++++++++++----
 requirements/docs.txt            |  2 ++
 4 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/docs-build.yml b/.github/workflows/docs-build.yml
index 13fa8deb5e..b630d1b237 100644
--- a/.github/workflows/docs-build.yml
+++ b/.github/workflows/docs-build.yml
@@ -31,15 +31,10 @@ jobs:
           python-version: "3.10"
 
       - name: Pull sphinx template
-        run: |
-          pip install -q "awscli >=1.30.0"
-          aws s3 sync --no-sign-request s3://sphinx-packages/ ${PYPI_LOCAL_DIR}
-          pip install lai-sphinx-theme -U -f ${PYPI_LOCAL_DIR}
+        run: make get-sphinx-theme
       - name: Install pandoc
         timeout-minutes: 5
-        run: |
-          sudo apt-get update --fix-missing
-          sudo apt-get install -y pandoc
+        run: sudo apt-get install -y pandoc
       - name: Install package & dependencies
         timeout-minutes: 20
         run: pip install . -U -r requirements/docs.txt
diff --git a/Makefile b/Makefile
index 2c27c12f5c..a50fded131 100644
--- a/Makefile
+++ b/Makefile
@@ -10,13 +10,13 @@ test: clean
 	python -m coverage run --source thunder -m pytest thunder tests -v
 	python -m coverage report
 
-sphinx-theme:
+get-sphinx-theme:
 	pip install -q awscli
 	mkdir -p dist/
 	aws s3 sync --no-sign-request s3://sphinx-packages/ dist/
 	pip install lai-sphinx-theme -f dist/
 
-docs: clean sphinx-theme
+docs: clean get-sphinx-theme
 	pip install -e . --quiet -r requirements/docs.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html
 	cd docs ; python -m sphinx -b html -W --keep-going source build
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index a13fede4b2..f11eeae4c3 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -17,7 +17,21 @@
 import sys
 from importlib.util import module_from_spec, spec_from_file_location
 
-import lai_sphinx_theme
+from lightning_utilities.core.imports import package_available
+
+_INSTALLED_NEW_THEME = package_available("lai_sphinx_theme")
+
+if _INSTALLED_NEW_THEME:
+    import lai_sphinx_theme
+else:
+    import pt_lightning_sphinx_theme
+    import warnings
+
+    warnings.warn(
+        "You are using the old theme, please install the new theme 'lai_sphinx_theme';"
+        " you can do this by running 'make get-sphinx-theme'."
+    )
+
 
 _PATH_HERE = os.path.abspath(os.path.dirname(__file__))
 _PATH_ROOT = os.path.realpath(os.path.join(_PATH_HERE, "..", ".."))
@@ -98,8 +112,9 @@ def _transform_changelog(path_in: str, path_out: str) -> None:
     "sphinx_copybutton",
     "sphinx_paramlinks",
     "sphinx_togglebutton",
-    "lai_sphinx_theme.extensions.lightning",
 ]
+if _INSTALLED_NEW_THEME:
+    extensions.append("lai_sphinx_theme.extensions.lightning")
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
@@ -152,8 +167,10 @@ def _transform_changelog(path_in: str, path_out: str) -> None:
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = "lai_sphinx_theme"
-html_theme_path = [lai_sphinx_theme.get_html_theme_path()]
+html_theme = "lai_sphinx_theme" if _INSTALLED_NEW_THEME else "pt_lightning_sphinx_theme"
+html_theme_path = [
+    lai_sphinx_theme.get_html_theme_path() if _INSTALLED_NEW_THEME else pt_lightning_sphinx_theme.get_html_theme_path()
+]
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
diff --git a/requirements/docs.txt b/requirements/docs.txt
index 64b473ccd8..3398d6cb4a 100644
--- a/requirements/docs.txt
+++ b/requirements/docs.txt
@@ -14,3 +14,5 @@ sphinx-copybutton ==0.5.2
 
 # installed from S3 location and fetched in advance
 lai-sphinx-theme
+# alternative /back-up (old) theme
+pt-lightning-sphinx-theme @ https://github.com/PyTorchLightning/lightning_sphinx_theme/archive/master.zip