tests/shell/test_sweeps.py

import pytest

from tests.helpers.run_command import run_command

"""
A couple of tests executing hydra sweeps.

Use the following command to skip slow tests:
    pytest -k "not slow"
"""


@pytest.mark.slow
def test_experiments():
    """Test running all available experiment configs for 1 epoch."""
    command = ["train.py", "-m", "experiment=glob(*)", "++trainer.max_epochs=1"]
    run_command(command)


@pytest.mark.slow
def test_default_sweep():
    """Test default Hydra sweeper."""
    command = [
        "train.py",
        "-m",
        "datamodule.batch_size=64,128",
        "model.lr=0.01,0.02",
        "trainer=default",
        "++trainer.fast_dev_run=true",
    ]
    run_command(command)


@pytest.mark.slow
def test_optuna_sweep():
    """Test Optuna sweeper."""
    command = [
        "train.py",
        "-m",
        "hparams_search=mnist_optuna",
        "trainer=default",
        "++trainer.fast_dev_run=true",
    ]
    run_command(command)