default test logger (Lightning-AI#1478)

* default test logger

* fix tests

* spawn

* try

* simplify tests

* simplify tests

* formatting

* loggers

* loggers

* revert to TestTube

* default

* default

* wraps

* world size

* optim imports

Author: Borda

update.sh renamed to .update.sh

@@ -102,10 +102,18 @@ def auto_add_sampler(self, dataloader: DataLoader, train: bool) -> DataLoader:
                 sampler = DistributedSampler(
                     dataloader.dataset,
                     num_replicas=xm.xrt_world_size(),
-                    rank=xm.get_ordinal()
+                    rank=xm.get_ordinal(),
                 )
             else:
-                sampler = DistributedSampler(dataloader.dataset)
+                world_size = {
+                    'ddp': self.num_nodes * self.num_processes,
+                    'ddp2': self.num_nodes,
+                }
+                sampler = DistributedSampler(
+                    dataloader.dataset,
+                    num_replicas=world_size.get(self.distributed_backend, 0),
+                    rank=self.proc_rank,
+                )
 
             dl_args['sampler'] = sampler
             dataloader = type(dataloader)(**dl_args)

pytorch_lightning/trainer/data_loading.py

@@ -23,11 +23,7 @@
 from pytorch_lightning.trainer.deprecated_api import TrainerDeprecatedAPITillVer0_8, TrainerDeprecatedAPITillVer0_9
 from pytorch_lightning.trainer.distrib_data_parallel import TrainerDDPMixin
 from pytorch_lightning.trainer.distrib_parts import (
-    TrainerDPMixin,
-    parse_gpu_ids,
-    determine_root_gpu_device,
-    pick_multiple_gpus,
-)
+    TrainerDPMixin, parse_gpu_ids, determine_root_gpu_device, pick_multiple_gpus)
 from pytorch_lightning.trainer.evaluation_loop import TrainerEvaluationLoopMixin
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.trainer.model_hooks import TrainerModelHooksMixin
@@ -736,13 +732,10 @@ def fit(
                 self.ddp_train(task, model)
             else:
                 self.__set_random_port()
-
                 # track for predict
                 self.model = model
-
                 # train
                 mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
-
                 # load weights if not interrupted
                 self.load_spawn_weights(model)
                 self.model = model

pytorch_lightning/trainer/trainer.py

@@ -1,3 +1,18 @@
 import os
 
+import numpy as np
+import torch
+
 TEST_ROOT = os.path.dirname(__file__)
+PACKAGE_ROOT = os.path.dirname(TEST_ROOT)
+TEMP_PATH = os.path.join(PACKAGE_ROOT, 'test_temp')
+
+# generate a list of random seeds for each test
+RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
+ROOT_SEED = 1234
+torch.manual_seed(ROOT_SEED)
+np.random.seed(ROOT_SEED)
+RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
+
+if not os.path.isdir(TEMP_PATH):
+    os.mkdir(TEMP_PATH)

tests/__init__.py

@@ -2,7 +2,6 @@
 
 import torch
 
-from tests.base.models import TestModelBase, DictHparamsModel
 from tests.base.eval_model_template import EvalModelTemplate
 from tests.base.mixins import (
     LightEmptyTestStep,
@@ -31,6 +30,7 @@
     LightTestNoneOptimizerMixin,
     LightZeroLenDataloader
 )
+from tests.base.models import TestModelBase, DictHparamsModel
 
 
 class LightningTestModel(LightTrainDataloader,

tests/base/__init__.py

@@ -7,10 +7,10 @@
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from tests import TEST_ROOT
+from tests import PACKAGE_ROOT
 
 #: local path to test datasets
-PATH_DATASETS = os.path.join(TEST_ROOT, 'Datasets')
+PATH_DATASETS = os.path.join(PACKAGE_ROOT, 'Datasets')
 
 
 class MNIST(Dataset):

tests/base/datasets.py

@@ -7,7 +7,7 @@
 
 
 # from test_models import assert_ok_test_acc, load_model, \
-#     clear_save_dir, get_default_testtube_logger, get_default_hparams, init_save_dir, \
+#     clear_save_dir, get_default_logger, get_default_hparams, init_save_dir, \
 #     init_checkpoint_callback, reset_seed, set_random_master_port
 
 

tests/base/debug.py

@@ -2,18 +2,18 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from tests.base.datasets import TrialMNIST
 from pytorch_lightning.core.lightning import LightningModule
+from tests.base.datasets import TrialMNIST
 from tests.base.eval_model_optimizers import ConfigureOptimizersPool
 from tests.base.eval_model_test_dataloaders import TestDataloaderVariations
 from tests.base.eval_model_test_epoch_ends import TestEpochEndVariations
 from tests.base.eval_model_test_steps import TestStepVariations
 from tests.base.eval_model_train_dataloaders import TrainDataloaderVariations
 from tests.base.eval_model_train_steps import TrainingStepVariations
+from tests.base.eval_model_utils import ModelTemplateUtils
 from tests.base.eval_model_valid_dataloaders import ValDataloaderVariations
 from tests.base.eval_model_valid_epoch_ends import ValidationEpochEndVariations
 from tests.base.eval_model_valid_steps import ValidationStepVariations
-from tests.base.eval_model_utils import ModelTemplateUtils
 
 
 class EvalModelTemplate(

tests/base/eval_model_template.py

@@ -1,4 +1,5 @@
 from torch.utils.data import DataLoader
+
 from tests.base.datasets import TrialMNIST
 
 

tests/base/eval_model_utils.py

@@ -1,5 +1,6 @@
 from abc import ABC
 from collections import OrderedDict
+
 import torch
 
 

tests/base/eval_model_valid_steps.py

@@ -1,4 +1,3 @@
-import os
 from collections import OrderedDict
 from typing import Dict
 

tests/base/models.py

@@ -7,18 +7,11 @@
 # from pl_examples import LightningTemplateModel
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.loggers import TestTubeLogger, TensorBoardLogger
-from tests.base import LightningTestModel, EvalModelTemplate
+from pytorch_lightning.loggers import TensorBoardLogger
+from tests import TEMP_PATH, RANDOM_PORTS, RANDOM_SEEDS
+from tests.base import LightningTestModel
 from tests.base.datasets import PATH_DATASETS
 
-# generate a list of random seeds for each test
-RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
-ROOT_SEED = 1234
-torch.manual_seed(ROOT_SEED)
-np.random.seed(ROOT_SEED)
-RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
-ROOT_PATH = os.path.abspath(os.path.dirname(__file__))
-
 
 def assert_speed_parity(pl_times, pt_times, num_epochs):
 
@@ -33,7 +26,7 @@ def assert_speed_parity(pl_times, pt_times, num_epochs):
         f"lightning was slower than PT (threshold {max_diff_per_epoch})"
 
 
-def run_model_test_no_loggers(trainer_options, model, min_acc=0.50):
+def run_model_test_without_loggers(trainer_options, model, min_acc=0.50):
     # save_dir = trainer_options['default_root_dir']
 
     # fit model
@@ -66,14 +59,16 @@ def run_model_test(trainer_options, model, on_gpu=True):
     save_dir = trainer_options['default_root_dir']
 
     # logger file to get meta
-    logger = get_default_testtube_logger(save_dir, False)
+    logger = get_default_logger(save_dir)
 
     # logger file to get weights
     checkpoint = init_checkpoint_callback(logger)
 
     # add these to the trainer options
-    trainer_options['checkpoint_callback'] = checkpoint
-    trainer_options['logger'] = logger
+    trainer_options.update(
+        checkpoint_callback=checkpoint,
+        logger=logger,
+    )
 
     # fit model
     trainer = Trainer(**trainer_options)
@@ -118,8 +113,10 @@ def get_default_hparams(continue_training=False, hpc_exp_number=0):
     }
 
     if continue_training:
-        args['test_tube_do_checkpoint_load'] = True
-        args['hpc_exp_number'] = hpc_exp_number
+        args.update(
+            test_tube_do_checkpoint_load=True,
+            hpc_exp_number=hpc_exp_number,
+        )
 
     hparams = Namespace(**args)
     return hparams
@@ -137,9 +134,9 @@ def get_default_model(lbfgs=False):
     return model, hparams
 
 
-def get_default_testtube_logger(save_dir, debug=True, version=None):
+def get_default_logger(save_dir, version=None):
     # set up logger object without actually saving logs
-    logger = TestTubeLogger(save_dir, name='lightning_logs', debug=debug, version=version)
+    logger = TensorBoardLogger(save_dir, name='lightning_logs', version=version)
     return logger
 
 
@@ -153,17 +150,20 @@ def get_data_path(expt_logger, path_dir=None):
         return expt.get_data_path(name, version)
     # the other experiments...
     if not path_dir:
-        path_dir = ROOT_PATH
+        if hasattr(expt_logger, 'save_dir') and expt_logger.save_dir:
+            path_dir = expt_logger.save_dir
+        else:
+            path_dir = TEMP_PATH
     path_expt = os.path.join(path_dir, name, 'version_%s' % version)
     # try if the new sub-folder exists, typical case for test-tube
     if not os.path.isdir(path_expt):
         path_expt = path_dir
     return path_expt
 
 
-def load_model(exp, root_weights_dir, module_class=LightningTestModel, path_expt=None):
+def load_model(logger, root_weights_dir, module_class=LightningTestModel, path_expt=None):
     # load trained model
-    path_expt_dir = get_data_path(exp, path_dir=path_expt)
+    path_expt_dir = get_data_path(logger, path_dir=path_expt)
     tags_path = os.path.join(path_expt_dir, TensorBoardLogger.NAME_CSV_TAGS)
 
     checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]

tests/base/utils.py

@@ -1,3 +1,5 @@
+from functools import wraps
+
 import pytest
 
 import torch.multiprocessing as mp
@@ -7,16 +9,12 @@ def pytest_configure(config):
     config.addinivalue_line("markers", "spawn: spawn test in a separate process using torch.multiprocessing.spawn")
 
 
-def wrap(i, fn, args):
-    return fn(*args)
-
-
 @pytest.mark.tryfirst
 def pytest_pyfunc_call(pyfuncitem):
     if pyfuncitem.get_closest_marker("spawn"):
         testfunction = pyfuncitem.obj
         funcargs = pyfuncitem.funcargs
         testargs = tuple([funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames])
 
-        mp.spawn(wrap, (testfunction, testargs))
+        mp.spawn(wraps, (testfunction, testargs))
         return True

tests/conftest.py

@@ -7,7 +7,6 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import (
     TensorBoardLogger, MLFlowLogger, NeptuneLogger, TestTubeLogger, CometLogger)
-from tests.base import LightningTestModel
 
 
 def _get_logger_args(logger_class, save_dir):