From 3ff5327e83caa0a842afcbe085c89aa6218cd1df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= <aedu.waelchli@gmail.com>
Date: Fri, 25 Sep 2020 16:00:02 +0200
Subject: [PATCH] Mocking loggers (part 1, wandb) (#3596)

* mocking for wandb

* remove wandb import in amp test

* mock loggers in sphinx

* check tests

* Update extra.txt

* setup

* dev

* min

* revert

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Jirka Borovec <jirka@pytorchlightning.ai>
---
 docs/source/conf.py                |  1 +
 docs/source/experiment_logging.rst |  2 +-
 pytorch_lightning/loggers/wandb.py | 13 ++++++-------
 requirements/extra.txt             |  3 ++-
 requirements/loggers.txt           |  6 ++++++
 setup.py                           |  3 ++-
 tests/models/test_amp.py           |  2 --
 7 files changed, 18 insertions(+), 12 deletions(-)
 create mode 100644 requirements/loggers.txt

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 58258af1fc..d95ea874d1 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -330,6 +330,7 @@ if SPHINX_MOCK_REQUIREMENTS:
     # mock also base packages when we are on RTD since we don't install them there
     MOCK_PACKAGES += package_list_from_file(os.path.join(PATH_ROOT, 'requirements/base.txt'))
     MOCK_PACKAGES += package_list_from_file(os.path.join(PATH_ROOT, 'requirements/extra.txt'))
+    MOCK_PACKAGES += package_list_from_file(os.path.join(PATH_ROOT, 'requirements/loggers.txt'))
 
 MOCK_MANUAL_PACKAGES = [
     'torchvision',
diff --git a/docs/source/experiment_logging.rst b/docs/source/experiment_logging.rst
index 0229514879..4ccad84ef2 100644
--- a/docs/source/experiment_logging.rst
+++ b/docs/source/experiment_logging.rst
@@ -192,7 +192,7 @@ First, install the package:
 
 Then configure the logger and pass it to the :class:`~pytorch_lightning.trainer.trainer.Trainer`:
 
-.. testcode::
+.. code-block:: python
 
     from pytorch_lightning.loggers import WandbLogger
     wandb_logger = WandbLogger(offline=True)
diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py
index 51ceafe614..4fe308ce28 100644
--- a/pytorch_lightning/loggers/wandb.py
+++ b/pytorch_lightning/loggers/wandb.py
@@ -25,11 +25,9 @@ import torch.nn as nn
 try:
     import wandb
     from wandb.wandb_run import Run
-    _WANDB_AVAILABLE = True
 except ImportError:  # pragma: no-cover
     wandb = None
     Run = None
-    _WANDB_AVAILABLE = False
 
 from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment
 from pytorch_lightning.utilities import rank_zero_only
@@ -57,10 +55,11 @@ class WandbLogger(LightningLoggerBase):
             :func:`wandb.init` can be passed as keyword arguments in this logger.
 
     Example:
-        >>> from pytorch_lightning.loggers import WandbLogger
-        >>> from pytorch_lightning import Trainer
-        >>> wandb_logger = WandbLogger()
-        >>> trainer = Trainer(logger=wandb_logger)
+
+        from pytorch_lightning.loggers import WandbLogger
+        from pytorch_lightning import Trainer
+        wandb_logger = WandbLogger()
+        trainer = Trainer(logger=wandb_logger)
 
     See Also:
         - `Tutorial <https://app.wandb.ai/cayush/pytorchlightning/reports/
@@ -82,7 +81,7 @@ class WandbLogger(LightningLoggerBase):
         experiment=None,
         **kwargs
     ):
-        if not _WANDB_AVAILABLE:
+        if wandb is None:
             raise ImportError('You want to use `wandb` logger which is not installed yet,'  # pragma: no-cover
                               ' install it with `pip install wandb`.')
         super().__init__()
diff --git a/requirements/extra.txt b/requirements/extra.txt
index 7ff0905ef5..f361a32f1d 100644
--- a/requirements/extra.txt
+++ b/requirements/extra.txt
@@ -1,10 +1,11 @@
 # extended list of package dependencies to reach full functionality
 
+# TODO: this shall be removed as we mock them in tests
 neptune-client>=0.4.109
 comet-ml>=3.1.12
 mlflow>=1.0.0
 test_tube>=0.7.5
-wandb>=0.8.21, <0.10.0  # higher version mess with system paths
+
 matplotlib>=3.1.1
 # no need to install with [pytorch] as pytorch is already installed and torchvision is required only for Horovod examples
 horovod>=0.19.2
diff --git a/requirements/loggers.txt b/requirements/loggers.txt
new file mode 100644
index 0000000000..3ec7b25db4
--- /dev/null
+++ b/requirements/loggers.txt
@@ -0,0 +1,6 @@
+# all supported loggers
+neptune-client>=0.4.109
+comet-ml>=3.1.12
+mlflow>=1.0.0
+test_tube>=0.7.5
+wandb>=0.8.21
\ No newline at end of file
diff --git a/setup.py b/setup.py
index a1dffe8907..0580772ca3 100755
--- a/setup.py
+++ b/setup.py
@@ -65,10 +65,11 @@ def load_long_description():
 extras = {
     # 'docs': load_requirements(file_name='docs.txt'),
     'examples': load_requirements(file_name='examples.txt'),
+    'loggers': load_requirements(file_name='loggers.txt'),
     'extra': load_requirements(file_name='extra.txt'),
     'test': load_requirements(file_name='test.txt')
 }
-extras['dev'] = extras['extra'] + extras['test']
+extras['dev'] = extras['extra'] + extras['loggers'] + extras['test']
 extras['all'] = extras['dev'] + extras['examples']  # + extras['docs']
 
 # https://packaging.python.org/discussions/install-requires-vs-requirements /
diff --git a/tests/models/test_amp.py b/tests/models/test_amp.py
index b1d92cf8e6..69c23bddb2 100644
--- a/tests/models/test_amp.py
+++ b/tests/models/test_amp.py
@@ -1,9 +1,7 @@
 import os
-from unittest.mock import MagicMock
 
 import pytest
 import torch
-import wandb
 
 import tests.base.develop_pipelines as tpipes
 import tests.base.develop_utils as tutils