Unable to import lightning on Google Colab

[adamjstewart]

Bug description

I'm trying to import lightning on Google Colab but it doesn't seem possible at the moment.

How to reproduce the bug

First thing I tried was:

import lightning

Not installed. No problem, let me install it:

%pip install lightning
import lightning

Nope, that fails too. Let's try ipython's reload feature:

%load_ext autoreload
%autoreload 2
%pip install lightning
import lightning

Nope, unable to reload. I've also played around with importlib.reload in a few different possible places but no luck so far.

Note that you have to click "Runtime > Disconnect and delete runtime" and "Reconnect" between each step.

Error messages and logs

Full stacktrace

<ipython-input-1-8f5cbb005247> in <module>
      1 get_ipython().run_line_magic('pip', 'install lightning')
----> 2 import lightning

14 frames
/usr/local/lib/python3.9/dist-packages/lightning/__init__.py in <module>
     25 from lightning.fabric.fabric import Fabric  # noqa: E402
     26 from lightning.fabric.utilities.seed import seed_everything  # noqa: E402
---> 27 from lightning.pytorch.callbacks import Callback  # noqa: E402
     28 from lightning.pytorch.core import LightningDataModule, LightningModule  # noqa: E402
     29 from lightning.pytorch.trainer import Trainer  # noqa: E402

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/__init__.py in <module>
     24 
     25 from lightning.fabric.utilities.seed import seed_everything  # noqa: E402
---> 26 from lightning.pytorch.callbacks import Callback  # noqa: E402
     27 from lightning.pytorch.core import LightningDataModule, LightningModule  # noqa: E402
     28 from lightning.pytorch.trainer import Trainer  # noqa: E402

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/callbacks/__init__.py in <module>
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
---> 14 from lightning.pytorch.callbacks.batch_size_finder import BatchSizeFinder
     15 from lightning.pytorch.callbacks.callback import Callback
     16 from lightning.pytorch.callbacks.checkpoint import Checkpoint

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/callbacks/batch_size_finder.py in <module>
     22 
     23 import lightning.pytorch as pl
---> 24 from lightning.pytorch.callbacks.callback import Callback
     25 from lightning.pytorch.tuner.batch_size_scaling import _scale_batch_size
     26 from lightning.pytorch.utilities.exceptions import _TunerExitException, MisconfigurationException

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/callbacks/callback.py in <module>
     23 
     24 import lightning.pytorch as pl
---> 25 from lightning.pytorch.utilities.types import STEP_OUTPUT
     26 
     27 

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/utilities/__init__.py in <module>
     18 from lightning.fabric.utilities import LightningEnum  # noqa: F401
     19 from lightning.fabric.utilities import move_data_to_device  # noqa: F401
---> 20 from lightning.pytorch.accelerators.hpu import _HPU_AVAILABLE  # noqa: F401
     21 from lightning.pytorch.utilities.combined_loader import CombinedLoader  # noqa: F401
     22 from lightning.pytorch.utilities.enums import GradClipAlgorithmType  # noqa: F401

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/accelerators/__init__.py in <module>
     14 from lightning.fabric.accelerators.registry import _AcceleratorRegistry, call_register_accelerators
     15 from lightning.pytorch.accelerators.accelerator import Accelerator  # noqa: F401
---> 16 from lightning.pytorch.accelerators.cpu import CPUAccelerator  # noqa: F401
     17 from lightning.pytorch.accelerators.cuda import CUDAAccelerator  # noqa: F401
     18 from lightning.pytorch.accelerators.hpu import HPUAccelerator  # noqa: F401

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/accelerators/cpu.py in <module>
     20 from lightning.fabric.utilities.types import _DEVICE
     21 from lightning.pytorch.accelerators.accelerator import Accelerator
---> 22 from lightning.pytorch.utilities.exceptions import MisconfigurationException
     23 
     24 

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/utilities/exceptions.py in <module>
     15 
     16 from lightning.fabric.utilities.exceptions import MisconfigurationException  # noqa: F401
---> 17 from lightning.pytorch.utilities.imports import _PYTHON_GREATER_EQUAL_3_11_0
     18 
     19 

/usr/local/lib/python3.9/dist-packages/lightning/pytorch/utilities/imports.py in <module>
     28 _OMEGACONF_AVAILABLE = package_available("omegaconf")
     29 _POPTORCH_AVAILABLE = package_available("poptorch")
---> 30 _RICH_AVAILABLE = package_available("rich") and compare_version("rich", operator.ge, "10.2.2")
     31 _TORCHVISION_AVAILABLE = RequirementCache("torchvision")
     32 _LIGHTNING_COLOSSALAI_AVAILABLE = RequirementCache("lightning-colossalai")

/usr/local/lib/python3.9/dist-packages/lightning_utilities/core/imports.py in compare_version(package, op, version, use_base_version)
     75         else:
     76             # try pkg_resources to infer version
---> 77             pkg_version = Version(pkg_resources.get_distribution(package).version)
     78     except TypeError:
     79         # this is mocked by Sphinx, so it should return True to generate all summaries

/usr/local/lib/python3.9/dist-packages/pkg_resources/__init__.py in get_distribution(dist)
    476         dist = Requirement.parse(dist)
    477     if isinstance(dist, Requirement):
--> 478         dist = get_provider(dist)
    479     if not isinstance(dist, Distribution):
    480         raise TypeError("Expected string, Requirement, or Distribution", dist)

/usr/local/lib/python3.9/dist-packages/pkg_resources/__init__.py in get_provider(moduleOrReq)
    352     """Return an IResourceProvider for the named module or requirement"""
    353     if isinstance(moduleOrReq, Requirement):
--> 354         return working_set.find(moduleOrReq) or require(str(moduleOrReq))[0]
    355     try:
    356         module = sys.modules[moduleOrReq]

/usr/local/lib/python3.9/dist-packages/pkg_resources/__init__.py in require(self, *requirements)
    907         included, even if they were already activated in this working set.
    908         """
--> 909         needed = self.resolve(parse_requirements(requirements))
    910 
    911         for dist in needed:

/usr/local/lib/python3.9/dist-packages/pkg_resources/__init__.py in resolve(self, requirements, env, installer, replace_conflicting, extras)
    798                 # Oops, the "best" so far conflicts with a dependency
    799                 dependent_req = required_by[req]
--> 800                 raise VersionConflict(dist, req).with_context(dependent_req)
    801 
    802             # push the new requirements onto the stack

ContextualVersionConflict: (Pygments 2.6.1 (/usr/local/lib/python3.9/dist-packages), Requirement.parse('pygments<3.0.0,>=2.13.0'), {'rich'})

Environment

Current environment

#- Lightning Component: all imports
#- PyTorch Lightning Version: 2.0.0
#- Lightning App Version: N/A
#- PyTorch Version: 1.13.1
#- Python version: 3.9.16
#- OS: Linux
#- CUDA/cuDNN version: 11.6
#- GPU models and configuration: Default Google Compute Engine backend
#- How you installed Lightning: pip
#- Running environment of LightningApp: cloud

More info

Here is what I believe to be the source of the issue.

lightning depends on rich, which depends on pygments<3.0.0,>=2.13.0. So %pip install lightning installs a newer version of pygments and uninstalls the pygments==2.6.1 that comes with the environment by default. However, Jupyter relies on pygments to render the notebook, so it has already been imported. For some reason, reload doesn't seem to work (possibly related to ipython/ipython#13347).

So far, I've found two possible workarounds, but no fixes. One is to run the notebook, let it crash, then "Restart runtime" so that pygments has already been upgraded and the notebook imports that instead. The other is to change the install to:

%pip install lightning pygments==2.6.1

This forces pip to use an older version of rich that still supports older pygments. The former workaround is unideal because unfamiliar users won't understand the error message, while the latter workaround is unideal because it only works on Google Colab and only until pygments gets updated in the future. I'm trying to write a tutorial that works anywhere without either of these constraints.

I'm honestly not sure what if anything Lightning can do to solve this issue. One possibility would be to not check the version of rich available and instead rely on pip to do the right thing. Another possibility would be to request that Google update the default version of pygments available in Colab, but that's out of your control and may take a while. Of course, if anyone can figure out why reload isn't working, that would be the most ideal solution.

cc @tchaton

[adamjstewart]

Reported to Google Colab too: googlecolab/colabtools#3514

[awaelchli]

Been struggling with this too, thanks for opening the issue @adamjstewart! Maybe we can temporarily require pygments<=2.6.1 in Lightning until this is updated in Colab. What do you think @Borda @carmocca ?

[carmocca]

Another workaround is to only try to import rich when it's requested to use the rich progress bar.

Using the RequirementCache and evaluating it lazily

[awaelchli]

I linked a pull request that fixes this issue temporarily.

[adamjstewart]

This was also fixed on the Google side: googlecolab/colabtools#3514

Your fix is probably more reliable though since the Google fix will be needed again the next time a newer version is required.

[awaelchli]

Thanks for confirming it!

[fractaldna22]

yet another sin of rich, and respositories which require it and keep a shitty module alive