Fix TPU CI (Lightning-AI#16613)

byronyi · Feb 2, 2023 · 6f7276b · 6f7276b
1 parent 710ca4b
commit 6f7276b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/dockers/base-xla/tpu_workflow_fabric.jsonnet b/dockers/base-xla/tpu_workflow_fabric.jsonnet
@@ -38,6 +38,10 @@ local tputests = base.BaseTest {
       fi
 
       echo "--- Install packages ---"
+      # set particular PyTorch version
+      for fpath in `ls requirements/**/*.txt`; do
+        python requirements/pytorch/adjust-versions.py $fpath {PYTORCH_VERSION};
+      done
       PACKAGE_NAME=fabric pip install .[dev]
       pip list
 
@@ -48,6 +52,8 @@ local tputests = base.BaseTest {
       export XRT_TPU_CONFIG="tpu_worker;0;${KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS:7}"
 
       echo "--- Sanity check TPU availability ---"
+      python -c "import torch_xla; print(torch_xla)"
+      python -c "from lightning_fabric.accelerators.tpu import _XLA_AVAILABLE; print(str(_XLA_AVAILABLE))"
       python -c "from lightning_fabric.accelerators import TPUAccelerator; assert TPUAccelerator.is_available()"
       echo "Sanity check passed!"
 

diff --git a/dockers/base-xla/tpu_workflow_pytorch.jsonnet b/dockers/base-xla/tpu_workflow_pytorch.jsonnet
@@ -38,16 +38,24 @@ local tputests = base.BaseTest {
       fi
 
       echo "--- Install packages ---"
+      # set particular PyTorch version
+      for fpath in `ls requirements/**/*.txt`; do
+        python requirements/pytorch/adjust-versions.py $fpath {PYTORCH_VERSION};
+      done
       PACKAGE_NAME=pytorch pip install .[dev]
       pip list
 
       pip install -q -r .actions/requirements.txt
-      python .actions/assistant.py copy_replace_imports --source_dir="./tests" --source_import="lightning.fabric" --target_import="lightning_fabric"
+      python .actions/assistant.py copy_replace_imports --source_dir="./tests" \
+          --source_import="lightning.fabric,lightning.pytorch" \
+          --target_import="lightning_fabric,pytorch_lightning"
 
       echo $KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS
       export XRT_TPU_CONFIG="tpu_worker;0;${KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS:7}"
 
       echo "--- Sanity check TPU availability ---"
+      python -c "import torch_xla; print(torch_xla)"
+      python -c "from pytorch_lightning.accelerators.tpu import _XLA_AVAILABLE; print(str(_XLA_AVAILABLE))"
       python -c "from pytorch_lightning.accelerators import TPUAccelerator; assert TPUAccelerator.is_available()"
       echo "Sanity check passed!"