Lazily import dependencies for NeptuneLogger (#18573)

Author: awaelchli

src/lightning/pytorch/loggers/neptune.py

@@ -15,15 +15,11 @@
 Neptune Logger
 --------------
 """
-__all__ = [
-    "NeptuneLogger",
-]
-
 import contextlib
 import logging
 import os
 from argparse import Namespace
-from typing import Any, Dict, Generator, List, Optional, Set, Union
+from typing import Any, Dict, Generator, List, Optional, Set, TYPE_CHECKING, Union
 
 from lightning_utilities.core.imports import RequirementCache
 from torch import Tensor
@@ -35,30 +31,15 @@
 from lightning.pytorch.utilities.model_summary import ModelSummary
 from lightning.pytorch.utilities.rank_zero import rank_zero_only
 
-# neptune is available with two names on PyPI : `neptune` and `neptune-client`
-_NEPTUNE_AVAILABLE = RequirementCache("neptune>=1.0")
-_NEPTUNE_CLIENT_AVAILABLE = RequirementCache("neptune-client")
-
-if _NEPTUNE_AVAILABLE:
-    # >1.0 package structure
-    import neptune
+if TYPE_CHECKING:
     from neptune import Run
     from neptune.handler import Handler
-    from neptune.types import File
-    from neptune.utils import stringify_unsupported
-elif _NEPTUNE_CLIENT_AVAILABLE:
-    # <1.0 package structure
-    import neptune.new as neptune
-    from neptune.new import Run
-    from neptune.new.handler import Handler
-    from neptune.new.types import File
-    from neptune.new.utils import stringify_unsupported
-else:
-    # needed for tests, mocks and function signatures
-    neptune, Run, Handler, File, stringify_unsupported = None, None, None, None, None
 
 log = logging.getLogger(__name__)
 
+# neptune is available with two names on PyPI : `neptune` and `neptune-client`
+_NEPTUNE_AVAILABLE = RequirementCache("neptune>=1.0")
+_NEPTUNE_CLIENT_AVAILABLE = RequirementCache("neptune-client")
 _INTEGRATION_VERSION_KEY = "source_code/integrations/pytorch-lightning"
 
 
@@ -258,6 +239,11 @@ def __init__(
         if self._run_instance is not None:
             self._retrieve_run_data()
 
+            if _NEPTUNE_AVAILABLE:
+                from neptune.handler import Handler
+            else:
+                from neptune.new.handler import Handler
+
             # make sure that we've log integration version for outside `Run` instances
             root_obj = self._run_instance
             if isinstance(root_obj, Handler):
@@ -266,6 +252,11 @@ def __init__(
             root_obj[_INTEGRATION_VERSION_KEY] = pl.__version__
 
     def _retrieve_run_data(self) -> None:
+        if _NEPTUNE_AVAILABLE:
+            from neptune.handler import Handler
+        else:
+            from neptune.new.handler import Handler
+
         assert self._run_instance is not None
         root_obj = self._run_instance
         if isinstance(root_obj, Handler):
@@ -317,6 +308,12 @@ def _verify_input_arguments(
         run: Optional[Union["Run", "Handler"]],
         neptune_run_kwargs: dict,
     ) -> None:
+        if _NEPTUNE_AVAILABLE:
+            from neptune import Run
+            from neptune.handler import Handler
+        else:
+            from neptune.new import Run
+            from neptune.new.handler import Handler
         # check if user passed the client `Run`/`Handler` object
         if run is not None and not isinstance(run, (Run, Handler)):
             raise ValueError("Run parameter expected to be of type `neptune.Run`, or `neptune.handler.Handler`.")
@@ -325,8 +322,8 @@ def _verify_input_arguments(
         any_neptune_init_arg_passed = any(arg is not None for arg in [api_key, project, name]) or neptune_run_kwargs
         if run is not None and any_neptune_init_arg_passed:
             raise ValueError(
-                "When an already initialized run object is provided"
-                " you can't provide other neptune.init_run() parameters.\n"
+                "When an already initialized run object is provided, you can't provide other `neptune.init_run()`"
+                " parameters."
             )
 
     def __getstate__(self) -> Dict[str, Any]:
@@ -336,12 +333,17 @@ def __getstate__(self) -> Dict[str, Any]:
         return state
 
     def __setstate__(self, state: Dict[str, Any]) -> None:
+        if _NEPTUNE_AVAILABLE:
+            import neptune
+        else:
+            import neptune.new as neptune
+
         self.__dict__ = state
         self._run_instance = neptune.init_run(**self._neptune_init_args)
 
     @property
     @rank_zero_experiment
-    def experiment(self) -> Run:
+    def experiment(self) -> "Run":
         r"""
         Actual Neptune run object. Allows you to use neptune logging features in your
         :class:`~lightning.pytorch.core.module.LightningModule`.
@@ -371,7 +373,12 @@ def training_step(self, batch, batch_idx):
 
     @property
     @rank_zero_experiment
-    def run(self) -> Run:
+    def run(self) -> "Run":
+        if _NEPTUNE_AVAILABLE:
+            import neptune
+        else:
+            import neptune.new as neptune
+
         if not self._run_instance:
             self._run_instance = neptune.init_run(**self._neptune_init_args)
             self._retrieve_run_data()
@@ -416,6 +423,11 @@ def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  #
             neptune_logger.log_hyperparams(PARAMS)
 
         """
+        if _NEPTUNE_AVAILABLE:
+            from neptune.utils import stringify_unsupported
+        else:
+            from neptune.new.utils import stringify_unsupported
+
         params = _convert_params(params)
         params = _sanitize_callable_params(params)
 
@@ -468,8 +480,13 @@ def save_dir(self) -> Optional[str]:
 
     @rank_zero_only
     def log_model_summary(self, model: "pl.LightningModule", max_depth: int = -1) -> None:
+        if _NEPTUNE_AVAILABLE:
+            from neptune.types import File
+        else:
+            from neptune.new.types import File
+
         model_str = str(ModelSummary(model=model, max_depth=max_depth))
-        self.run[self._construct_path_with_prefix("model/summary")] = neptune.types.File.from_content(
+        self.run[self._construct_path_with_prefix("model/summary")] = File.from_content(
             content=model_str, extension="txt"
         )
 
@@ -484,6 +501,11 @@ def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
         if not self._log_model_checkpoints:
             return
 
+        if _NEPTUNE_AVAILABLE:
+            from neptune.types import File
+        else:
+            from neptune.new.types import File
+
         file_names = set()
         checkpoints_namespace = self._construct_path_with_prefix("model/checkpoints")
 

tests/tests_pytorch/loggers/conftest.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import sys
 from types import ModuleType
-from unittest.mock import Mock
+from unittest.mock import MagicMock, Mock
 
 import pytest
 
@@ -37,6 +37,8 @@ def mlflow_mock(monkeypatch):
 
     mlflow.tracking = mlflow_tracking
     mlflow.entities = mlflow_entities
+
+    monkeypatch.setattr("lightning.pytorch.loggers.mlflow._MLFLOW_AVAILABLE", True),
     return mlflow
 
 
@@ -73,6 +75,8 @@ class RunType:  # to make isinstance checks pass
     wandb.sdk = wandb_sdk
     wandb.sdk.lib = wandb_sdk_lib
     wandb.wandb_run = wandb_wandb_run
+
+    monkeypatch.setattr("lightning.pytorch.loggers.wandb._WANDB_AVAILABLE", True)
     return wandb
 
 
@@ -92,4 +96,46 @@ def comet_mock(monkeypatch):
     monkeypatch.setitem(sys.modules, "comet_ml.api", comet_api)
 
     comet.api = comet_api
+
+    monkeypatch.setattr("lightning.pytorch.loggers.comet._COMET_AVAILABLE", True)
     return comet
+
+
+@pytest.fixture()
+def neptune_mock(monkeypatch):
+    class RunType:  # to make isinstance checks pass
+        def get_root_object(self):
+            pass
+
+        def __getitem__(self, item):
+            pass
+
+        def __setitem__(self, key, value):
+            pass
+
+    run_mock = MagicMock(spec=RunType, exists=Mock(return_value=False), wait=Mock(), get_structure=MagicMock())
+    run_mock.get_root_object.return_value = run_mock
+
+    neptune = ModuleType("neptune")
+    neptune.init_run = Mock(return_value=run_mock)
+    neptune.Run = RunType
+    monkeypatch.setitem(sys.modules, "neptune", neptune)
+
+    neptune_handler = ModuleType("handler")
+    neptune_handler.Handler = RunType
+    monkeypatch.setitem(sys.modules, "neptune.handler", neptune_handler)
+
+    neptune_types = ModuleType("types")
+    neptune_types.File = Mock()
+    monkeypatch.setitem(sys.modules, "neptune.types", neptune_types)
+
+    neptune_utils = ModuleType("utils")
+    neptune_utils.stringify_unsupported = Mock()
+    monkeypatch.setitem(sys.modules, "neptune.utils", neptune_utils)
+
+    neptune.handler = neptune_handler
+    neptune.types = neptune_types
+    neptune.utils = neptune_utils
+
+    monkeypatch.setattr("lightning.pytorch.loggers.neptune._NEPTUNE_AVAILABLE", True)
+    return neptune

tests/tests_pytorch/loggers/test_all.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import contextlib
 import inspect
 import os
 import pickle
@@ -37,17 +36,7 @@
 from tests_pytorch.helpers.runif import RunIf
 from tests_pytorch.loggers.test_comet import _patch_comet_atexit
 from tests_pytorch.loggers.test_mlflow import mock_mlflow_run_creation
-from tests_pytorch.loggers.test_neptune import create_neptune_mock
-
-LOGGER_CTX_MANAGERS = (
-    mock.patch("lightning.pytorch.loggers.mlflow._MLFLOW_AVAILABLE", return_value=True),
-    mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock()),
-    mock.patch("lightning.pytorch.loggers.neptune.neptune", new_callable=create_neptune_mock),
-    mock.patch("lightning.pytorch.loggers.neptune._NEPTUNE_AVAILABLE", return_value=True),
-    mock.patch("lightning.pytorch.loggers.neptune.Run", new=mock.Mock),
-    mock.patch("lightning.pytorch.loggers.neptune.Handler", new=mock.Mock),
-    mock.patch("lightning.pytorch.loggers.neptune.File", new=mock.Mock()),
-)
+
 ALL_LOGGER_CLASSES = (
     CometLogger,
     CSVLogger,
@@ -79,18 +68,11 @@ def _instantiate_logger(logger_class, save_dir, **override_kwargs):
 
 
 @mock.patch.dict(os.environ, {})
-@mock.patch("lightning.pytorch.loggers.wandb._WANDB_AVAILABLE", True)
-@mock.patch("lightning.pytorch.loggers.comet._COMET_AVAILABLE", True)
+@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
 @pytest.mark.parametrize("logger_class", ALL_LOGGER_CLASSES)
-def test_loggers_fit_test_all(logger_class, mlflow_mock, wandb_mock, comet_mock, tmp_path):
+def test_loggers_fit_test_all(logger_class, mlflow_mock, wandb_mock, comet_mock, neptune_mock, tmp_path):
     """Verify that basic functionality of all loggers."""
-    with contextlib.ExitStack() as stack:
-        for mgr in LOGGER_CTX_MANAGERS:
-            stack.enter_context(mgr)
-        _test_loggers_fit_test(tmp_path, logger_class)
-
 
-def _test_loggers_fit_test(tmp_path, logger_class):
     class CustomModel(BoringModel):
         def training_step(self, batch, batch_idx):
             loss = self.step(batch)
@@ -300,38 +282,32 @@ def _test_logger_initialization(tmp_path, logger_class):
 
 
 @mock.patch.dict(os.environ, {})
-def test_logger_with_prefix_all(mlflow_mock, wandb_mock, comet_mock, monkeypatch, tmp_path):
+@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
+def test_logger_with_prefix_all(mlflow_mock, wandb_mock, comet_mock, neptune_mock, monkeypatch, tmp_path):
     """Test that prefix is added at the beginning of the metric keys."""
     prefix = "tmp"
 
     # Comet
-    with mock.patch("lightning.pytorch.loggers.comet._COMET_AVAILABLE", return_value=True):
-        _patch_comet_atexit(monkeypatch)
-        logger = _instantiate_logger(CometLogger, save_dir=tmp_path, prefix=prefix)
-        logger.log_metrics({"test": 1.0}, step=0)
-        logger.experiment.log_metrics.assert_called_once_with({"tmp-test": 1.0}, epoch=None, step=0)
+    _patch_comet_atexit(monkeypatch)
+    logger = _instantiate_logger(CometLogger, save_dir=tmp_path, prefix=prefix)
+    logger.log_metrics({"test": 1.0}, step=0)
+    logger.experiment.log_metrics.assert_called_once_with({"tmp-test": 1.0}, epoch=None, step=0)
 
     # MLflow
-    with mock.patch("lightning.pytorch.loggers.mlflow._MLFLOW_AVAILABLE", return_value=True), mock.patch(
-        "lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock()
-    ):
-        Metric = mlflow_mock.entities.Metric
-        logger = _instantiate_logger(MLFlowLogger, save_dir=tmp_path, prefix=prefix)
-        logger.log_metrics({"test": 1.0}, step=0)
-        logger.experiment.log_batch.assert_called_once_with(
-            run_id=ANY, metrics=[Metric(key="tmp-test", value=1.0, timestamp=ANY, step=0)]
-        )
+    Metric = mlflow_mock.entities.Metric
+    logger = _instantiate_logger(MLFlowLogger, save_dir=tmp_path, prefix=prefix)
+    logger.log_metrics({"test": 1.0}, step=0)
+    logger.experiment.log_batch.assert_called_once_with(
+        run_id=ANY, metrics=[Metric(key="tmp-test", value=1.0, timestamp=ANY, step=0)]
+    )
 
     # Neptune
-    with mock.patch("lightning.pytorch.loggers.neptune.neptune"), mock.patch(
-        "lightning.pytorch.loggers.neptune._NEPTUNE_AVAILABLE", return_value=True
-    ), mock.patch("lightning.pytorch.loggers.neptune.Handler", new=mock.Mock):
-        logger = _instantiate_logger(NeptuneLogger, api_key="test", project="project", save_dir=tmp_path, prefix=prefix)
-        assert logger.experiment.__getitem__.call_count == 0
-        logger.log_metrics({"test": 1.0}, step=0)
-        assert logger.experiment.__getitem__.call_count == 1
-        logger.experiment.__getitem__.assert_called_with("tmp/test")
-        logger.experiment.__getitem__().append.assert_called_once_with(1.0)
+    logger = _instantiate_logger(NeptuneLogger, api_key="test", project="project", save_dir=tmp_path, prefix=prefix)
+    assert logger.experiment.__getitem__.call_count == 0
+    logger.log_metrics({"test": 1.0}, step=0)
+    assert logger.experiment.__getitem__.call_count == 1
+    logger.experiment.__getitem__.assert_called_with("tmp/test")
+    logger.experiment.__getitem__().append.assert_called_once_with(1.0)
 
     # TensorBoard
     if _TENSORBOARD_AVAILABLE:
@@ -345,14 +321,14 @@ def test_logger_with_prefix_all(mlflow_mock, wandb_mock, comet_mock, monkeypatch
     logger.experiment.add_scalar.assert_called_once_with("tmp-test", 1.0, 0)
 
     # WandB
-    with mock.patch("lightning.pytorch.loggers.wandb._WANDB_AVAILABLE", True):
-        logger = _instantiate_logger(WandbLogger, save_dir=tmp_path, prefix=prefix)
-        wandb_mock.run = None
-        wandb_mock.init().step = 0
-        logger.log_metrics({"test": 1.0}, step=0)
-        logger.experiment.log.assert_called_once_with({"tmp-test": 1.0, "trainer/global_step": 0})
+    logger = _instantiate_logger(WandbLogger, save_dir=tmp_path, prefix=prefix)
+    wandb_mock.run = None
+    wandb_mock.init().step = 0
+    logger.log_metrics({"test": 1.0}, step=0)
+    logger.experiment.log.assert_called_once_with({"tmp-test": 1.0, "trainer/global_step": 0})
 
 
+@mock.patch("lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock())
 def test_logger_default_name(mlflow_mock, monkeypatch, tmp_path):
     """Test that the default logger name is lightning_logs."""
     # CSV
@@ -370,14 +346,11 @@ def test_logger_default_name(mlflow_mock, monkeypatch, tmp_path):
     assert logger.name == "lightning_logs"
 
     # MLflow
-    with mock.patch("lightning.pytorch.loggers.mlflow._MLFLOW_AVAILABLE", return_value=True), mock.patch(
-        "lightning.pytorch.loggers.mlflow._get_resolve_tags", Mock()
-    ):
-        client = mlflow_mock.tracking.MlflowClient()
-        client.get_experiment_by_name.return_value = None
-        logger = _instantiate_logger(MLFlowLogger, save_dir=tmp_path)
-
-        _ = logger.experiment
-        logger._mlflow_client.create_experiment.assert_called_with(name="lightning_logs", artifact_location=ANY)
-        # on MLFLowLogger `name` refers to the experiment id
-        # assert logger.experiment.get_experiment(logger.name).name == "lightning_logs"
+    client = mlflow_mock.tracking.MlflowClient()
+    client.get_experiment_by_name.return_value = None
+    logger = _instantiate_logger(MLFlowLogger, save_dir=tmp_path)
+
+    _ = logger.experiment
+    logger._mlflow_client.create_experiment.assert_called_with(name="lightning_logs", artifact_location=ANY)
+    # on MLFLowLogger `name` refers to the experiment id
+    # assert logger.experiment.get_experiment(logger.name).name == "lightning_logs"