From e04d4fd25e66b2f5502bb4371d16aa8554bf878e Mon Sep 17 00:00:00 2001
From: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Date: Fri, 24 Mar 2023 09:46:45 +0100
Subject: [PATCH] ci: separate integrations (#17170)

* integrations

* Rename

* Apply suggestions from code review

* cut out

* skip _

* fix

* setup

* 1e-2

* uninstall

* revert

* note
---
 .actions/assistant.py                     | 2 +-
 .azure/gpu-tests-pytorch.yml              | 8 ++++++--
 .github/workflows/ci-tests-pytorch.yml    | 7 ++++++-
 requirements/_integrations/strategies.txt | 5 +++++
 requirements/pytorch/strategies.txt       | 2 --
 src/lightning/__setup__.py                | 2 +-
 6 files changed, 19 insertions(+), 7 deletions(-)
 create mode 100644 requirements/_integrations/strategies.txt

diff --git a/.actions/assistant.py b/.actions/assistant.py
index ca30b13dec..d791793af9 100644
--- a/.actions/assistant.py
+++ b/.actions/assistant.py
@@ -240,7 +240,7 @@ def _load_aggregate_requirements(req_dir: str = "requirements", freeze_requireme
         load_requirements(d, unfreeze="none" if freeze_requirements else "major")
         for d in glob.glob(os.path.join(req_dir, "*"))
         # skip empty folder as git artefacts, and resolving Will's special issue
-        if os.path.isdir(d) and len(glob.glob(os.path.join(d, "*"))) > 0 and "__pycache__" not in d
+        if os.path.isdir(d) and len(glob.glob(os.path.join(d, "*"))) > 0 and not os.path.basename(d).startswith("_")
     ]
     if not requires:
         return
diff --git a/.azure/gpu-tests-pytorch.yml b/.azure/gpu-tests-pytorch.yml
index 02b628e9c3..ceffccf8b6 100644
--- a/.azure/gpu-tests-pytorch.yml
+++ b/.azure/gpu-tests-pytorch.yml
@@ -111,13 +111,17 @@ jobs:
     - bash: pip install -e .[extra,test,examples] -U --find-links ${TORCH_URL}
       displayName: 'Install package & dependencies'
 
-    - bash: pip uninstall -y -r requirements/pytorch/strategies.txt
+    - bash: |
+        pip uninstall -y -r requirements/pytorch/strategies.txt \
+                         -r requirements/_integrations/strategies.txt
       condition: ne(variables['scope'], 'strategies')
       displayName: 'Uninstall strategies'
 
     - bash: |
         set -e
-        pip install -r requirements/pytorch/strategies.txt --find-links ${TORCH_URL}
+        pip install -r requirements/pytorch/strategies.txt \
+                    -r requirements/_integrations/strategies.txt \
+          --find-links ${TORCH_URL}
         python requirements/pytorch/check-avail-strategies.py
       condition: eq(variables['scope'], 'strategies')
       displayName: 'Install strategies'
diff --git a/.github/workflows/ci-tests-pytorch.yml b/.github/workflows/ci-tests-pytorch.yml
index 316710ffe1..666afd72e7 100644
--- a/.github/workflows/ci-tests-pytorch.yml
+++ b/.github/workflows/ci-tests-pytorch.yml
@@ -115,7 +115,9 @@ jobs:
 
     - name: Install package & dependencies
       run: |
-        pip install -e .[extra,test] "pytest-timeout" -U -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE} --prefer-binary
+        pip install -q pip -U
+        pip install .[extra,test] -U "pytest-timeout" \
+          -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE} --prefer-binary
         pip list
     - name: Dump handy wheels
       if: github.event_name == 'push' && github.ref == 'refs/heads/master'
@@ -145,6 +147,9 @@ jobs:
           --source_import="lightning.fabric,lightning.pytorch" \
           --target_import="lightning_fabric,pytorch_lightning"
 
+    - name: Prevent using raw source
+      run: rm -rf src/
+
     - name: Testing Warnings
       working-directory: tests/tests_pytorch
       # needs to run outside of `pytest`
diff --git a/requirements/_integrations/strategies.txt b/requirements/_integrations/strategies.txt
new file mode 100644
index 0000000000..3227549ab7
--- /dev/null
+++ b/requirements/_integrations/strategies.txt
@@ -0,0 +1,5 @@
+# NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
+#  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
+
+lightning-colossalai >=0.1.0
+lightning-bagua >=0.1.0
diff --git a/requirements/pytorch/strategies.txt b/requirements/pytorch/strategies.txt
index 3588468672..c8a5c9531f 100644
--- a/requirements/pytorch/strategies.txt
+++ b/requirements/pytorch/strategies.txt
@@ -2,5 +2,3 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 deepspeed>=0.6.0, <0.8.0  # TODO: Include 0.8.x after https://github.com/microsoft/DeepSpeed/commit/b587c7e85470329ac25df7c7c2521ff9b2833db7 gets released
-lightning-colossalai>=0.1.0dev1
-lightning-bagua>=0.1.0rc2
diff --git a/src/lightning/__setup__.py b/src/lightning/__setup__.py
index ceef349a96..15d4a47808 100644
--- a/src/lightning/__setup__.py
+++ b/src/lightning/__setup__.py
@@ -36,7 +36,7 @@ def _prepare_extras() -> Dict[str, Any]:
     extras = {
         f"{p.parent.name}-{p.stem}": _ASSISTANT.load_requirements(file_name=p.name, path_dir=p.parent, **common_args)
         for p in req_files
-        if p.name not in ("docs.txt", "devel.txt", "base.txt")
+        if p.name not in ("docs.txt", "devel.txt", "base.txt") and not p.parts[-2].startswith("_")
     }
     for extra in list(extras):
         name = "-".join(extra.split("-")[1:])