Fix XLAEnvironment detection on TPU pod #16806

awaelchli · 2023-02-18T14:15:57Z

What does this PR do?

The newly added test reproduces the issue and fails without the fix:

____________________ test_accelerator_choice_tpu[tpu-None] _____________________

accelerator = 'tpu', devices = None

    @RunIf(tpu=True, standalone=True)
    @pytest.mark.parametrize(
        ["accelerator", "devices"], [("tpu", None), ("tpu", 1), ("tpu", [1]), ("tpu", 8), ("auto", 1), ("auto", 8)]
    )
    @mock.patch.dict(os.environ, os.environ.copy(), clear=True)
    def test_accelerator_choice_tpu(accelerator, devices):
        connector = _Connector(accelerator=accelerator, devices=devices)
        assert isinstance(connector.accelerator, TPUAccelerator)
        if devices is None or (isinstance(devices, int) and devices > 1):
            # accelerator=tpu, devices=None (default) maps to devices=auto (8) and then chooses XLAStrategy
            # This behavior may change in the future: https://github.com/Lightning-AI/lightning/issues/10606
            assert isinstance(connector.strategy, XLAStrategy)
>           assert isinstance(connector.cluster_environment, XLAEnvironment)
E           assert False
E            +  where False = isinstance(<lightning_fabric.plugins.environments.lightning.LightningEnvironment object at 0x7fcb048549d0>, XLAEnvironment)
E            +    where <lightning_fabric.plugins.environments.lightning.LightningEnvironment object at 0x7fcb048549d0> = <lightning_fabric.connector._Connector object at 0x7fcb0482d070>.cluster_environment

test_connector.py:74: AssertionError
=========================== short test summary info ============================
FAILED test_connector.py::test_accelerator_choice_tpu[tpu-None] - assert False
======================== 1 failed, 8 warnings in 6.41s =========================

cc @Borda @carmocca @justusschock @awaelchli

github-actions · 2023-02-18T14:16:36Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 1.11)	success	✅
pl-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
pl-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 1.11)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.12)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 1.11)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.12)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 1.13)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.13)	success	✅
pl-cpu (windows-2022, pytorch, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py, src/lightning/fabric/connector.py.

🟢 pytorch_lightning: Azure HPU

Check ID	Status
pytorch-lightning (HPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py.

🟢 pytorch_lightning: Azure IPU

Check ID	Status
pytorch-lightning (IPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py.

🟢 pytorch_lightning: Docs

Check ID	Status
make-doctest (pytorch)	success	✅
make-html (pytorch)	success	✅

These checks are required after the changes to src/lightning/pytorch/trainer/connectors/accelerator_connector.py.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, lightning, 3.8, 1.11)	success	✅
fabric-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
fabric-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.9, 1.11)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 1.12)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.9, 1.11)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 1.12)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (macOS-11, fabric, 3.8, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 1.13)	success	✅
fabric-cpu (windows-2022, fabric, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, tests/tests_fabric/test_connector.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, tests/tests_fabric/test_connector.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py.

🟢 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.8)	success	✅
install-pkg (ubuntu-22.04, app, 3.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.10)	success	✅
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.10)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.10)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py.

🟢 link-check

Check ID	Status
markdown-link-check	success	✅

These checks are required after the changes to src/lightning/fabric/CHANGELOG.md, src/lightning/pytorch/CHANGELOG.md.

Thank you for your contribution! 💜

Note
This comment is automatically generated and updates for 60 minutes every 180 seconds. If you have any other questions, contact carmocca for help.

codecov · 2023-02-18T14:31:34Z

Codecov Report

Merging #16806 (8c59ba4) into master (6950a07) will decrease coverage by 22%.
The diff coverage is 100%.

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #16806     +/-   ##
=========================================
- Coverage      82%      59%    -22%     
=========================================
  Files         437      412     -25     
  Lines       31583    31290    -293     
=========================================
- Hits        25784    18543   -7241     
- Misses       5799    12747   +6948

for more information, see https://pre-commit.ci

…' into bugfix/xla-environment-detection

for more information, see https://pre-commit.ci

tpu fixes

awaelchli added 2 commits February 18, 2023 15:08

fix

d53b183

test

a8abb3a

awaelchli requested review from carmocca and justusschock as code owners February 18, 2023 14:15

awaelchli added the bug Something isn't working label Feb 18, 2023

github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Feb 18, 2023

awaelchli force-pushed the bugfix/xla-environment-detection branch from d391548 to a8abb3a Compare February 18, 2023 16:33

awaelchli requested a review from williamFalcon as a code owner February 18, 2023 16:33

pre-commit-ci bot and others added 5 commits February 18, 2023 16:34

[pre-commit.ci] auto fixes from pre-commit.com hooks

adf6fee

for more information, see https://pre-commit.ci

changelog

4ef2052

update tests

24d92c7

make sure both attributes are in sync

058c521

Merge remote-tracking branch 'origin/bugfix/xla-environment-detection…

a6f213a

…' into bugfix/xla-environment-detection

awaelchli added this to the v1.9.x milestone Feb 18, 2023

awaelchli added the fun Staff contributions outside working hours - to differentiate from the "community" label label Feb 18, 2023

awaelchli and others added 4 commits February 18, 2023 19:05

debug

8baf970

refactor test

5e21585

[pre-commit.ci] auto fixes from pre-commit.com hooks

f56e1db

for more information, see https://pre-commit.ci

avoid leaking env vars

8c59ba4

awaelchli force-pushed the bugfix/xla-environment-detection branch from a7716d3 to 8c59ba4 Compare February 19, 2023 19:01

awaelchli changed the title ~~WIP: Fix XLAEnvironment detection on TPU pod~~ Fix XLAEnvironment detection on TPU pod Feb 19, 2023

justusschock approved these changes Feb 20, 2023

View reviewed changes

carmocca approved these changes Feb 20, 2023

View reviewed changes

carmocca merged commit 2844e9e into master Feb 20, 2023

carmocca deleted the bugfix/xla-environment-detection branch February 20, 2023 10:01

mergify bot added the ready PRs ready to be merged label Feb 20, 2023

awaelchli added a commit that referenced this pull request Feb 20, 2023

Fix XLAEnvironment detection on TPU pod (#16806)

95254f3

awaelchli added a commit that referenced this pull request Feb 20, 2023

Fix XLAEnvironment detection on TPU pod (#16806)

6657052

awaelchli added a commit that referenced this pull request Feb 21, 2023

Fix XLAEnvironment detection on TPU pod (#16806)

61e6fd5

tpu fixes

carmocca added a commit that referenced this pull request Feb 21, 2023

Fix XLAEnvironment detection on TPU pod (#16806)

bdaba90

tpu fixes

lexierule pushed a commit that referenced this pull request Feb 21, 2023

Fix XLAEnvironment detection on TPU pod (#16806)

41a5c2a

tpu fixes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix XLAEnvironment detection on TPU pod #16806

Fix XLAEnvironment detection on TPU pod #16806

Uh oh!

awaelchli commented Feb 18, 2023 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Feb 18, 2023 •

edited

Loading

Uh oh!

codecov bot commented Feb 18, 2023 •

edited

Loading

Uh oh!

Uh oh!

Fix XLAEnvironment detection on TPU pod #16806

Fix XLAEnvironment detection on TPU pod #16806

Uh oh!

Conversation

awaelchli commented Feb 18, 2023 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

github-actions bot commented Feb 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

⚡ Required checks status: All passing 🟢

Groups summary

Uh oh!

codecov bot commented Feb 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

awaelchli commented Feb 18, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Feb 18, 2023 •

edited

Loading

codecov bot commented Feb 18, 2023 •

edited

Loading