ci: timeout & DATASETS_PATH (#1030)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.azure/gpu-tests.yml

@@ -32,8 +32,10 @@ jobs:
       DEVICES: $( python -c 'name = "$(Agent.Name)" ; gpus = name.split("_")[-1] if "_" in name else "0,1"; print(gpus)' )
       DEBIAN_FRONTEND: "noninteractive"
       TZ: "Europe/Amsterdam"
+      PIP_CACHE_DIR: "/var/tmp/pip"
       TORCH_HOME: "/var/tmp/torch"
       TRANSFORMERS_CACHE: "/var/tmp/huggingface"
+      DATASETS_PATH: "/var/tmp/bolts_datasets"  # custom for running Bolts
       FREEZE_REQUIREMENTS: 1
 
     container:
@@ -111,7 +113,9 @@ jobs:
       displayName: Download ROMs
 
     - bash: |
-        python -m pytest tests/ -v --cov=pl_bolts --junitxml=$(Build.StagingDirectory)/test-results.xml --durations=30
+        # todo: optimize runtime, mainly because of retina example
+        python -m pytest tests/ -v --cov=pl_bolts --durations=30 \
+          --junitxml=$(Build.StagingDirectory)/test-results.xml
       displayName: 'Testing'
 
     - bash: |

.github/workflows/ci-tests.yml

@@ -117,7 +117,7 @@ jobs:
       run: python -m atari_py.import_roms ROMS
 
     - name: Tests
-      run: python -m pytest tests -v --cov=pl_bolts
+      run: python -m pytest tests -v --cov=pl_bolts --timeout=200
 
     - name: Statistics
       if: success()

requirements/test.txt

@@ -1,6 +1,7 @@
 coverage[toml] >7.0.0, <8.0.0
 pytest ==7.3.2
 pytest-cov ==4.1.0
+pytest-timeout ==2.1.0
 pytest-rerunfailures ==11.1.2
 
 scikit-learn >=1.0.2, <=1.2.2

tests/__init__.py

@@ -5,7 +5,7 @@
 
 TEST_ROOT = os.path.realpath(os.path.dirname(__file__))
 PROJECT_ROOT = os.path.dirname(TEST_ROOT)
-DATASETS_PATH = os.path.join(PROJECT_ROOT, "_datasets")
+DATASETS_PATH = os.environ.get("DATASETS_PATH", os.path.join(PROJECT_ROOT, "_datasets"))
 # generate a list of random seeds for each test
 ROOT_SEED = 1234
 

tests/callbacks/verification/test_batch_gradient.py

@@ -115,10 +115,7 @@ def test_batch_gradient_verification_pl_module(mix_data, device):
     assert result == is_valid
 
 
-@pytest.mark.parametrize(
-    "gpus",
-    [0, pytest.param(1, marks=pytest.mark.skipif(**_MARK_REQUIRE_GPU))],
-)
+@pytest.mark.parametrize("gpus", [0, pytest.param(1, marks=pytest.mark.skipif(**_MARK_REQUIRE_GPU))])
 def test_batch_gradient_verification_callback(gpus):
     """Test detection of batch gradient mixing with the callback implementation."""
     trainer = Trainer(gpus=gpus)

tests/models/rl/integration/test_actor_critic_models.py

@@ -1,23 +1,21 @@
 import argparse
 
+import torch.cuda
 from pl_bolts.models.rl.advantage_actor_critic_model import AdvantageActorCritic
 from pl_bolts.models.rl.sac_model import SAC
 from pytorch_lightning import Trainer
 
 
-def test_a2c():
+def test_a2c_cli():
     """Smoke test that the A2C model runs."""
 
     parent_parser = argparse.ArgumentParser(add_help=False)
     parent_parser = AdvantageActorCritic.add_model_specific_args(parent_parser)
-    args_list = [
-        "--env",
-        "CartPole-v0",
-    ]
+    args_list = ["--env", "CartPole-v0"]
     hparams = parent_parser.parse_args(args_list)
 
     trainer = Trainer(
-        gpus=0,
+        gpus=int(torch.cuda.is_available()),
         max_steps=100,
         max_epochs=100,  # Set this as the same as max steps to ensure that it doesn't stop early
         val_check_interval=1,  # This just needs 'some' value, does not effect training right now
@@ -27,7 +25,7 @@ def test_a2c():
     trainer.fit(model)
 
 
-def test_sac():
+def test_sac_cli():
     """Smoke test that the SAC model runs."""
 
     parent_parser = argparse.ArgumentParser(add_help=False)
@@ -37,7 +35,7 @@ def test_sac():
         "--warm_start_size",
         "100",
         "--gpus",
-        "0",
+        "0",  # fixme: RuntimeError on GPU: Expected all tensors to be on the same device
         "--env",
         "Pendulum-v0",
         "--batch_size",

tests/models/rl/integration/test_policy_models.py

@@ -1,6 +1,7 @@
 import argparse
 from unittest import TestCase
 
+import torch
 from pl_bolts.models.rl.reinforce_model import Reinforce
 from pl_bolts.models.rl.vanilla_policy_gradient_model import VanillaPolicyGradient
 from pytorch_lightning import Trainer
@@ -10,14 +11,11 @@ class TestPolicyModels(TestCase):
     def setUp(self) -> None:
         parent_parser = argparse.ArgumentParser(add_help=False)
         parent_parser = VanillaPolicyGradient.add_model_specific_args(parent_parser)
-        args_list = [
-            "--env",
-            "CartPole-v0",
-        ]
+        args_list = ["--env", "CartPole-v0"]
         self.hparams = parent_parser.parse_args(args_list)
 
         self.trainer = Trainer(
-            gpus=0,
+            gpus=int(torch.cuda.is_available()),
             max_steps=100,
             max_epochs=100,  # Set this as the same as max steps to ensure that it doesn't stop early
             val_check_interval=1,  # This just needs 'some' value, does not effect training right now

tests/models/rl/integration/test_value_models.py

@@ -2,6 +2,7 @@
 from unittest import TestCase
 
 import pytest
+import torch
 from pl_bolts.models.rl.double_dqn_model import DoubleDQN
 from pl_bolts.models.rl.dqn_model import DQN
 from pl_bolts.models.rl.dueling_dqn_model import DuelingDQN
@@ -19,7 +20,7 @@ def setUp(self) -> None:
             "--warm_start_size",
             "100",
             "--gpus",
-            "0",
+            str(int(torch.cuda.is_available())),
             "--env",
             "PongNoFrameskip-v4",
         ]

tests/models/rl/unit/test_a2c.py

@@ -9,12 +9,7 @@ def test_a2c_loss():
     """Test the reinforce loss function."""
     parent_parser = argparse.ArgumentParser(add_help=False)
     parent_parser = AdvantageActorCritic.add_model_specific_args(parent_parser)
-    args_list = [
-        "--env",
-        "CartPole-v0",
-        "--batch_size",
-        "32",
-    ]
+    args_list = ["--env", "CartPole-v0", "--batch_size", "32"]
     hparams = parent_parser.parse_args(args_list)
     model = AdvantageActorCritic(**vars(hparams))
 
@@ -31,12 +26,7 @@ def test_a2c_train_batch():
     """Tests that a single batch generates correctly."""
     parent_parser = argparse.ArgumentParser(add_help=False)
     parent_parser = AdvantageActorCritic.add_model_specific_args(parent_parser)
-    args_list = [
-        "--env",
-        "CartPole-v0",
-        "--batch_size",
-        "32",
-    ]
+    args_list = ["--env", "CartPole-v0", "--batch_size", "32"]
     hparams = parent_parser.parse_args(args_list)
     model = AdvantageActorCritic(**vars(hparams))
 

tests/models/rl/unit/test_reinforce.py

@@ -23,14 +23,7 @@ def setUp(self) -> None:
 
         parent_parser = argparse.ArgumentParser(add_help=False)
         parent_parser = Reinforce.add_model_specific_args(parent_parser)
-        args_list = [
-            "--env",
-            "CartPole-v0",
-            "--batch_size",
-            "32",
-            "--gamma",
-            "0.99",
-        ]
+        args_list = ["--env", "CartPole-v0", "--batch_size", "32", "--gamma", "0.99"]
         self.hparams = parent_parser.parse_args(args_list)
         self.model = Reinforce(**vars(self.hparams))
 

tests/models/rl/unit/test_sac.py

@@ -9,12 +9,7 @@ def test_sac_loss():
     """Test the reinforce loss function."""
     parent_parser = argparse.ArgumentParser(add_help=False)
     parent_parser = SAC.add_model_specific_args(parent_parser)
-    args_list = [
-        "--env",
-        "Pendulum-v0",
-        "--batch_size",
-        "32",
-    ]
+    args_list = ["--env", "Pendulum-v0", "--batch_size", "32"]
     hparams = parent_parser.parse_args(args_list)
     model = SAC(**vars(hparams))
 
@@ -36,12 +31,7 @@ def test_sac_train_batch():
     """Tests that a single batch generates correctly."""
     parent_parser = argparse.ArgumentParser(add_help=False)
     parent_parser = SAC.add_model_specific_args(parent_parser)
-    args_list = [
-        "--env",
-        "Pendulum-v0",
-        "--batch_size",
-        "32",
-    ]
+    args_list = ["--env", "Pendulum-v0", "--batch_size", "32"]
     hparams = parent_parser.parse_args(args_list)
     model = SAC(**vars(hparams))
 

tests/models/rl/unit/test_vpg.py

@@ -20,12 +20,7 @@ def setUp(self) -> None:
 
         parent_parser = argparse.ArgumentParser(add_help=False)
         parent_parser = VanillaPolicyGradient.add_model_specific_args(parent_parser)
-        args_list = [
-            "--env",
-            "CartPole-v0",
-            "--batch_size",
-            "32",
-        ]
+        args_list = ["--env", "CartPole-v0", "--batch_size", "32"]
         self.hparams = parent_parser.parse_args(args_list)
         self.model = VanillaPolicyGradient(**vars(self.hparams))
 

tests/models/self_supervised/test_ssl_scripts.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import pytest
+import torch
 
 from tests import _MARK_REQUIRE_GPU, DATASETS_PATH
 
@@ -9,12 +10,7 @@
 
 # todo: failing for GPU as some is on CPU other on GPU
 @pytest.mark.skip(reason="FIXME: failing for GPU as some is on CPU other on GPU")
-@pytest.mark.parametrize(
-    "cli_args",
-    [
-        _DEFAULT_ARGS + " --gpus 1",
-    ],
-)
+@pytest.mark.parametrize("cli_args", [_DEFAULT_ARGS + f" --gpus {int(torch.cuda.is_available())}"])
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 def test_cli_run_ssl_amdim(cli_args):
     """Test running CLI for an example with default params."""
@@ -31,7 +27,8 @@ def test_cli_run_ssl_amdim(cli_args):
 @pytest.mark.parametrize(
     "cli_args",
     [
-        _DEFAULT_ARGS + " --dataset %(dataset_name)s" " --hidden_mlp 512" " --encoder resnet18" " --gpus 1",
+        _DEFAULT_ARGS
+        + f" --dataset %(dataset_name)s --hidden_mlp 512 --encoder resnet18 --gpus {int(torch.cuda.is_available())}"
     ],
 )
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
@@ -44,12 +41,7 @@ def test_cli_run_ssl_cpc(cli_args, dataset_name):
         cli_main()
 
 
-@pytest.mark.parametrize(
-    "cli_args",
-    [
-        _DEFAULT_ARGS + " --gpus 1",
-    ],
-)
+@pytest.mark.parametrize("cli_args", [_DEFAULT_ARGS + f" --gpus {int(torch.cuda.is_available())}"])
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 def test_cli_run_ssl_moco(cli_args):
     """Test running CLI for an example with default params."""
@@ -60,12 +52,7 @@ def test_cli_run_ssl_moco(cli_args):
         cli_main()
 
 
-@pytest.mark.parametrize(
-    "cli_args",
-    [
-        _DEFAULT_ARGS + " --online_ft" " --gpus 1" " --fp32",
-    ],
-)
+@pytest.mark.parametrize("cli_args", [_DEFAULT_ARGS + f" --online_ft --gpus {int(torch.cuda.is_available())} --fp32"])
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 def test_cli_run_ssl_simclr(cli_args):
     """Test running CLI for an example with default params."""
@@ -76,12 +63,7 @@ def test_cli_run_ssl_simclr(cli_args):
         cli_main()
 
 
-@pytest.mark.parametrize(
-    "cli_args",
-    [
-        _DEFAULT_ARGS + " --gpus 1",
-    ],
-)
+@pytest.mark.parametrize("cli_args", [_DEFAULT_ARGS + f" --gpus {int(torch.cuda.is_available())}"])
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 def test_cli_run_ssl_byol(cli_args):
     """Test running CLI for an example with default params."""
@@ -102,7 +84,7 @@ def test_cli_run_ssl_byol(cli_args):
         " --sinkhorn_iterations 1"
         " --nmb_prototypes 2"
         " --queue_length 0"
-        " --gpus 1",
+        f" --gpus {int(torch.cuda.is_available())}"
     ],
 )
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
@@ -115,12 +97,7 @@ def test_cli_run_ssl_swav(cli_args):
         cli_main()
 
 
-@pytest.mark.parametrize(
-    "cli_args",
-    [
-        _DEFAULT_ARGS + " --gpus 1",
-    ],
-)
+@pytest.mark.parametrize("cli_args", [_DEFAULT_ARGS + f" --gpus {int(torch.cuda.is_available())}"])
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 def test_cli_run_ssl_simsiam(cli_args):
     """Test running CLI for an example with default params."""

tests/models/test_scripts.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import pytest
+import torch
 from pl_bolts.utils import _JSONARGPARSE_GREATER_THAN_4_16_0
 
 from tests import _MARK_REQUIRE_GPU, DATASETS_PATH
@@ -109,7 +110,9 @@ def test_cli_run_vision_image_gpt(cli_args):
         cli_main()
 
 
-@pytest.mark.parametrize("cli_args", [_DEFAULT_LIGHTNING_CLI_ARGS + " --trainer.gpus 1"])
+@pytest.mark.parametrize(
+    "cli_args", [_DEFAULT_LIGHTNING_CLI_ARGS + f" --trainer.gpus {int(torch.cuda.is_available())}"]
+)
 @pytest.mark.skipif(**_MARK_REQUIRE_GPU)
 # FixMe; see https://github.com/omni-us/jsonargparse/issues/187
 @pytest.mark.skipif(not _JSONARGPARSE_GREATER_THAN_4_16_0, reason="Failing on CI, need to be fixed")