docs: dont mock imports when running sphinx doctest (Lightning-AI#2396)

* skip if no amp

* dont mock when doctesting

* install extra

Author: awaelchli

.github/workflows/docs-checks.yml

@@ -40,11 +40,14 @@ jobs:
         run: |
           # python -m pip install --upgrade --user pip
           pip install -r requirements/base.txt -U -f https://download.pytorch.org/whl/torch_stable.html -q
+          pip install -r requirements/extra.txt
           pip install -r requirements/docs.txt
           python --version ; pip --version ; pip list
         shell: bash
 
       - name: Test Documentation
+        env:
+          SPHINX_MOCK_REQUIREMENTS: 0
         run: |
           # First run the same pipeline as Read-The-Docs
           apt-get update && sudo apt-get install -y cmake

docs/source/apex.rst

@@ -21,6 +21,7 @@ Native torch
 When using PyTorch 1.6+ Lightning uses the native amp implementation to support 16-bit.
 
 .. testcode::
+    :skipif: not APEX_AVAILABLE and not NATIVE_AMP_AVALAIBLE
 
     # turn on 16-bit
     trainer = Trainer(precision=16)
@@ -62,6 +63,7 @@ Enable 16-bit
 ^^^^^^^^^^^^^
 
 .. testcode::
+    :skipif: not APEX_AVAILABLE and not NATIVE_AMP_AVALAIBLE
 
     # turn on 16-bit
     trainer = Trainer(amp_level='O2', precision=16)
@@ -76,6 +78,7 @@ TPU 16-bit
 16-bit on TPus is much simpler. To use 16-bit with TPUs set precision to 16 when using the tpu flag
 
 .. testcode::
+    :skipif: not XLA_AVAILABLE
 
     # DEFAULT
     trainer = Trainer(tpu_cores=8, precision=32)

docs/source/conf.py

@@ -416,7 +416,11 @@ def find_source():
 import os
 import torch
 
-TORCHVISION_AVAILABLE = importlib.util.find_spec('torchvision')
+from pytorch_lightning.utilities import NATIVE_AMP_AVALAIBLE
+APEX_AVAILABLE = importlib.util.find_spec("apex") is not None
+XLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None
+TORCHVISION_AVAILABLE = importlib.util.find_spec("torchvision") is not None
+
 
 """
 coverage_skip_undoc_in_source = True

pytorch_lightning/trainer/__init__.py

@@ -721,6 +721,7 @@ def on_train_end(self, trainer, pl_module):
 will still show torch.float32.
 
 .. testcode::
+    :skipif: not APEX_AVAILABLE and not NATIVE_AMP_AVALAIBLE
 
     # default used by the Trainer
     trainer = Trainer(precision=32)