Minor miscellaneous fixes

.github/checkgroup.yml

@@ -143,6 +143,7 @@ subprojects:
       - "build-cuda (3.9, 1.12, 11.7.1)"
       - "build-cuda (3.9, 1.13, 12.0.1)"
       - "build-cuda (3.10, 2.0, 12.0.1)"
+      - "build-cuda (3.10, 2.0, 11.7.1)"
       #- "build-NGC"
       - "build-pl (3.9, 1.11, 11.3.1)"
       - "build-pl (3.9, 1.12, 11.7.1)"

.github/workflows/ci-dockers.yml

@@ -74,6 +74,8 @@ jobs:
           - {python_version: "3.9", pytorch_version: "1.12", cuda_version: "11.7.1"}
           - {python_version: "3.9", pytorch_version: "1.13", cuda_version: "12.0.1"}
           - {python_version: "3.10", pytorch_version: "2.0", cuda_version: "12.0.1"}
+          # these are used in Azure GPU CI
+          - {python_version: "3.10", pytorch_version: "2.0", cuda_version: "11.7.1"}
     steps:
       - uses: actions/checkout@v3
       - uses: docker/setup-buildx-action@v2

.github/workflows/ci-tests-fabric.yml

@@ -112,8 +112,7 @@ jobs:
       run: |
         python -m pip install -q pip -U
         extra=$(python -c "print({'lightning': 'fabric-'}.get('${{ matrix.pkg-name }}', ''))")
-        pip install -e ".[${extra}test]" "pytest-timeout" -U -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE_DIR} --prefer-binary
-        pip install -r requirements/fabric/strategies.txt -f ${PYPI_CACHE_DIR} --prefer-binary
+        pip install -e ".[${extra}test,${extra}strategies]" "pytest-timeout" -U -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE_DIR} --prefer-binary
         pip list
     - name: Dump handy wheels
       if: github.event_name == 'push' && github.ref == 'refs/heads/master'

.github/workflows/ci-tests-pytorch.yml

@@ -116,7 +116,7 @@ jobs:
       run: |
         python -m pip install -q pip -U
         extra=$(python -c "print({'lightning': 'pytorch-'}.get('${{ matrix.pkg-name }}', ''))")
-        pip install ".[${extra}extra,${extra}test]" -U \
+        pip install ".[${extra}extra,${extra}test,${extra}strategies]" -U \
           "pytest-timeout" -r requirements/_integrations/accelerators.txt \
           -f ${TORCH_URL} ${TORCH_PREINSTALL} -f ${PYPI_CACHE_DIR} --prefer-binary
         pip list

.github/workflows/tpu-tests.yml

@@ -115,7 +115,7 @@ jobs:
       env:
         JOB_NAME: ${{ env.PR_NUMBER }}-${{ matrix.pkg-name }}-${{ matrix.runtime }}-${{ env.SHA }}
       run: |
-        gcloud compute tpus tpu-vm create "$JOB_NAME" --accelerator-type=v4-8 --version="tpu-vm-v4-pt-$XLA_VER"
+        gcloud compute tpus tpu-vm create "$JOB_NAME" --accelerator-type=v4-8 --version="tpu-vm-v4-pt-$XLA_VER" --preemptible
 
     - name: Cancel job
       if: steps.tpu-create.outcome != 'success'

requirements/fabric/base.txt

@@ -4,6 +4,6 @@
 numpy >=1.17.2, <1.25.1
 torch >=1.11.0, <2.1.0
 fsspec[http]>2021.06.0, <2023.5.0
-packaging >=17.1, <=23.0
+packaging >=20.0, <=23.0
 typing-extensions >=4.0.0, <=4.4.0
 lightning-utilities >=0.8.0, <0.10.0

requirements/pytorch/base.txt

@@ -7,6 +7,6 @@ tqdm >=4.57.0, <4.66.0
 PyYAML >=5.4, <=6.0
 fsspec[http] >2021.06.0, <2023.5.0
 torchmetrics >=0.7.0, <1.1.0  # needed for using fixed compare_version
-packaging >=17.1, <=23.0
+packaging >=20.0, <=23.0
 typing-extensions >=4.0.0, <=4.4.0
 lightning-utilities >=0.8.0, <0.10.0

tests/tests_fabric/strategies/test_fsdp_integration.py

@@ -321,7 +321,13 @@ def test_setup_with_orig_params_and_multiple_param_groups():
 
 @RunIf(min_cuda_gpus=2, skip_windows=True, standalone=True, dynamo=True)
 @mock.patch.dict(os.environ, {})
-@pytest.mark.parametrize("compile_after_setup", [False, True])
+@pytest.mark.parametrize(
+    "compile_after_setup",
+    [
+        False,
+        pytest.param(True, marks=RunIf(min_python="3.9", reason="https://github.com/pytorch/pytorch/issues/97811")),
+    ],
+)
 def test_compile(compile_after_setup):
     """Test that the model can be compiled before and after the model is wrapped in FSDP."""
     model = BoringModel()

tests/tests_pytorch/conftest.py

@@ -155,7 +155,7 @@ def __new__(cls, self, *args, **kwargs):
 
         # torch doesn't allow creation of mps devices on older versions
         monkeypatch.setattr("torch.device", MpsDeviceMock)
-    monkeypatch.setattr(lightning.fabric.accelerators.mps, "_get_all_available_mps_gpus", lambda: list(range(n)))
+    monkeypatch.setattr(lightning.fabric.accelerators.mps, "_get_all_available_mps_gpus", lambda: [0] if n > 0 else [])
     monkeypatch.setattr(lightning.fabric.accelerators.mps.MPSAccelerator, "is_available", lambda *_: n > 0)
 
 
@@ -169,16 +169,6 @@ def mps_count_1(monkeypatch):
     mock_mps_count(monkeypatch, 1)
 
 
-@pytest.fixture()
-def mps_count_2(monkeypatch):
-    mock_mps_count(monkeypatch, 2)
-
-
-@pytest.fixture()
-def mps_count_4(monkeypatch):
-    mock_mps_count(monkeypatch, 4)
-
-
 def mock_xla_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> None:
     monkeypatch.setattr(lightning.pytorch.strategies.xla, "_XLA_AVAILABLE", value)
     monkeypatch.setattr(lightning.pytorch.strategies.single_xla, "_XLA_AVAILABLE", value)

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -781,6 +781,7 @@ def get_defaults(cls):
 
 
 @RunIf(min_cuda_gpus=1)  # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite
+@pytest.mark.xfail(raises=ImportError, reason="Not updated to latest API")
 @pytest.mark.skipif(not package_available("lightning_colossalai"), reason="Requires Colossal AI Strategy")
 def test_colossalai_external_strategy(monkeypatch):
     with mock.patch(
@@ -795,6 +796,7 @@ def test_colossalai_external_strategy(monkeypatch):
 
 
 @RunIf(min_cuda_gpus=1)  # trigger this test on our GPU pipeline, because we don't install the package on the CPU suite
+@pytest.mark.xfail(raises=ImportError, reason="Not updated to latest API")
 @pytest.mark.skipif(not package_available("lightning_bagua"), reason="Requires Bagua Strategy")
 def test_bagua_external_strategy(monkeypatch):
     with mock.patch(